AtCoder Beginner Contest 220 E,F問題メモ

AtCoder Beginner Contest 220

E - Distance on Large Perfect Binary Tree

• $2^N-1$ 個の頂点からなる完全二分木
  • 頂点 $1$ を根とし、頂点 $i$ の子は $2i,2i+1$
• 頂点 $(i,j)$ の組で、その最短距離が $D$ となるものの個数を $\mod{998244353}$ で求めよ
• $1 \le N \le 10^6$

パターンが同じになるものを見つけて、まとめて計算する。

$D=1$ の場合、答えは辺数×2 $=2^{N+1}-4$ となる。

以下、$D \ge 2$ とする。

完全二分木なので、深さが同じ頂点以下の部分木は形状が全く一緒となる。
当然、その中だけで取れる答えの対象ペアの個数も一緒となる。

一方、深さが異なる頂点の部分木同士はまとめて考えることはできないため、深さでグループ分けするのがよさそう。

深さ $d$ に位置する頂点（$v$ とする）毎に、 「$v$ が2頂点を結ぶパスに含まれ、かつそれが最も浅い頂点である」組の個数を数えると、重複せずに数え上げられる。
以下の2つに分けて考える。

• A.$v$を一端、もう一端をそこより $D$ 深い頂点とする組
• B.$v$ の左の子孫から1点、右の子孫から1点選び、距離が $D$ となる2頂点の組

N=5                               d
               ①                 0
        ,-----'  `-----,
       ②              ③         1
     ／  ＼          ／  ＼
   ④      ⑤      ⑥      ⑦     2
  ／＼    ／＼    ／＼    ／＼
 ⑧  ⑨  ⑩  ⑪  ⑫  ⑬  ⑭  ⑮   3
 /\  /\  /\  /\  /\  /\  /\  /\
⑯⑰⑱⑲⑳㉑㉒㉓㉔㉕㉖㉗㉘㉙㉚㉛  4

上記の例で、例えば $D=3$ のとき、②にとっての (②,⑯), (②,㉒) などが該当する。

$d+D \ge N$ なら、$D$ だけ深い頂点は存在しない。

$d+D \lt N$ の場合、1頂点につき $2^D$ のペア対象が存在する。（②に対する⑯～㉓）

今回の問題では $(i,j),(j,i)$ は別に数えるので2倍、さらに深さ $d$ には $2^d$ 個の頂点があるので、

全部で $2^{d+D+1}$ となる。

上記の例で、例えば $D=3$ のとき、②にとっての (⑧,⑤), (④,⑪) などが該当する。

左の頂点を深さ $d+l$ の頂点とすると、距離 $D$ にするためには右は深さ $d+(D-l)$ となる。

②にとって、自身より $l$ 深い左の頂点は $2^{l-1}$ 個、$D-l$ 深い右の頂点は $2^{D-l-1}$ 個で、 この組合せは $2^{D-2}$ 個となる。

つまり、$l$ に関わらず、どの深さとどの深さをペアにしても組合せは常に $2^{D-2}$ で固定である。いい感じ。

では何通りの $l$ が取れるのかだが、

• 両方を最大の深さにしても、$D$ に足りない場合
  • 1つも取れないのでスキップ

• 最小の $1$ から最大の $D-1$ までフルで取れる場合
  • $D-1$ 通り
• $l$ が小さすぎると相方が深すぎて存在しなくなる場合
  • 相方の深さがMAX $d+D-l = N-1$ のとき、$l$ が取れる最小値 $d+D-N+1$
  • 従って $d+D-N+1～N-1$ の $2N-d-D-1$ 通り

この2つの小さい方が、$l$ が取れるパターン数（$p$ とする）となる。

A.と同様、反転させたペアで2倍、深さ $d$ には $2^d$ 個の頂点があるので、

全部で $p \times 2^{d+D-1}$ となる。

$d=0～N-2$ まで足し合わせた結果が答え。

Python3

n, d = map(int, input().split())
MOD = 998244353

if d == 1:
    ans = (pow(2, n + 1, MOD) - 4) % MOD
    print(ans)
    exit()

pow2 = [0] * max(n + 5, d + 5)
pow2[0] = 1
for i in range(1, max(n + 5, d + 5)):
    pow2[i] = pow2[i - 1] * 2 % MOD

ans = 0
for dep in range(n - 1):
    dd = n - dep - 1  # 下端までの距離
    if 2 * dd < d:
        break

    if dd >= d:
        ans += pow2[d + 1] * pow2[dep]

    if d > 1:
        ans += pow2[d - 1] * min(d - 1, 2 * dd - d + 1) * pow2[dep]
        ans %= MOD

print(ans)

F - Distance Sums 2

• $N$ 頂点の木、$i$ 番目の辺は $u_i,v_i$ を結ぶ無向辺
• 各頂点について、他の全頂点との距離の総和を求めよ
  • つまり、$i=1,2,...,N$ について、$\displaystyle \sum_{j=1}^{N}dist(i,j)$ を求めよ
• $2 \le N \le 2 \times 10^5$

主客転倒で考える。

1本の辺が、どの頂点にどれだけ寄与するのか考える。

 ... ○--,  e     ,--○
      ○-○---○-○-○  ...
   ○----'     `--○
└────┘ └──────┘
     A              B

こうなったとき、辺 $e$ は、

• $A$ 側の頂点にとって、$B$ 側のどれか1個とペアになったときに1回ずつ、計 $|B|$ 回答えに加算される
• $B$ 側の頂点にとって、$A$ 側のどれか1個とペアになったときに1回ずつ、計 $|A|$ 回答えに加算される

つまり、$A$ 側の頂点に全て $|B|$ を加算し、$B$ 側の頂点に全て $|A|$ を加算すればよい。
これを全ての辺に対して行えば、答えとなる。

加算には、木の累積和を用いる。

• AtCoder ABC 187 E - Through Path (水色, 500 点) - けんちょんの競プロ精進記録

適当に頂点1を根としておいて、まず1回目のDFSで各頂点の部分木のサイズ $size[v]$ と深さを求める。

次いで、以下の配列を用意する。最初は全て0。

• $Acc[v]=v$ 以下の部分木の全頂点に、$Acc[v]$ の値を加算する

各辺について、両端の頂点の内、深さが浅い方（親）を $u$、深い方（子）を $v$ とすると、

• $v$ 側の全頂点に、$u$ 側のサイズ $N-size[v]$ を加算する
  • $Acc[v]+=N-size[v]$

• $u$ 側の全頂点に、$v$ 側のサイズ $size[v]$ を加算する
  • これは、まず全頂点に $size[v]$ を加算し、$v$ 以下の頂点ではそれを打ち消せばよい
  • $Acc[1]+=size[v]$
  • $Acc[v]-=size[v]$

最後に2回目のDFSで、$Acc$ を根から順に子に伝播させていけばよい。

Python3

n = int(input())
edges = []
links = [set() for _ in range(n)]
for _ in range(n - 1):
    u, v = map(int, input().split())
    u -= 1
    v -= 1

    edges.append((u, v))
    links[u].add(v)
    links[v].add(u)

depth = [-1] * n
subtree_size = [-1] * n
depth[0] = 0
q = [0]
while q:
    v = q[-1]
    if subtree_size[v] == -1:
        subtree_size[v] = -2
        for u in links[v]:
            if subtree_size[u] != -1:
                continue
            depth[u] = depth[v] + 1
            q.append(u)
    else:
        q.pop()
        sts = 1
        for u in links[v]:
            if subtree_size[u] < 0:
                continue
            sts += subtree_size[u]
        subtree_size[v] = sts

acc = [0] * n
for u, v in edges:
    if depth[u] > depth[v]:
        u, v = v, u

    size = subtree_size[v]
    acc[v] += n - 2 * size
    acc[0] += size

q = [0]
while q:
    v = q.pop()
    for u in links[v]:
        if depth[u] < depth[v]:
            continue
        acc[u] += acc[v]
        q.append(u)

print(*acc)