AtCoder Beginner Contest 376（Promotion of AtCoder Career Design DAY）G問題メモ

AtCoder Beginner Contest 376（Promotion of AtCoder Career Design DAY）

G - Treasure Hunting

問題文

頂点に $0$ から $N$ までの番号がついた $N + 1$ 頂点の根付き木があります。
頂点 $1$ , 頂点 $2$ , $\dots$ , 頂点 $N$ のうちどこか $1$ 頂点に宝が隠されています。
- 頂点 $i$ に宝がある確率は $\frac{a_i}{\sum_{j=1}^N a_j}$ です。
また、各頂点には「探索済」と「未探索」のどちらか一方の状態を持ちます。
はじめ頂点 $0$ は探索済で、それ以外の頂点は未探索です。
あなたは、宝がある頂点が探索済になるまで以下の操作を行います。
- 親が探索済であるような未探索の頂点を選び、その頂点を探索済にする。
操作回数の期待値が最小になるように行動した時の操作回数の期待値を $\mod{998244353}$ で求めてください。
$T$ 個のテストケースが与えられるので、それぞれに対して答えを求めてください。

制約

$1 \leq T \leq 2 \times 10^5$
$1 \leq N \leq 2 \times 10^5$
$0 \leq p_i \lt i$
$1 \leq a_i$
$\sum_{i=1}^N a_i \leq 10^8$
全てのテストケースに対する $N$ の総和は $2 \times 10^5$ 以下
入力される値は全て整数

解法

選ぶ頂点の順番 $q=(q_1,...,q_N)$ を決めたとき、期待値は、各頂点に宝がある確率に操作回数をかけて、

$\displaystyle \sum_{i=1}^{N}i\frac{a_{qi}}{S} = \frac{1}{S}\sum_{i=1}^{N}ia_{qi}$
- ただし $S=\sum a_i$

となる。定数を除いた $\displaystyle \sum_{i=1}^{N}ia_{qi}$ を「コスト」とし、コストを最小化するような $q$ を特定することを目指す。
$q$ 内では、 $p_i$ より先に $i$ を置いてはいけないという制約を満たす必要がある。

根（優先的に選ぶ頂点）から、または葉（最後まで残す頂点）から、何らかの指標で貪欲に決めていけないか？と考えるが、選んだ枝を途中で保留して他の枝に移っても良いことなどもあり、よい指標が決められない。

この問題では、 $q$ をジグソーパズルのように「断片的に結合していく」解法がとれる。

取っかかりとして、以下の事実がある。

$a_i$ 最大の頂点は、選べる状態になったら即座に選んで損しない。

略証
① ...→○→ p → ... → j → i → ○ → ...
② ...→○→ p → ... → i → j → ○ → ...
      ----------------|          |------     共通

i が ai 最大の頂点、p を i の親とする。
i が操作順で p の直後でないとき、1つ前に操作したのを j として、
この2数だけに着目したコストは、操作を k, k+1 番目として

① k * aj + (k+1) * ai
② k * ai + (k+1) * aj

これは明らかに②の方が小さい。また、どの頂点の選び順の制約にも影響を及ぼさない。
よって、①のような状況は②に置き換えても損せず、
これを繰り返すと、「i は p の直後が最適」となる。

場所はともかく、q には ...,(p,i),... という連続した列が含まれることが分かる。

これを列同士の結合に一般化する。

上記の $(p,i)$ のように、断片的に「この並びは連続する」ことがわかっているものを「列」とする。
「列 $b$ の後は、すぐに列 $c$ が来る」ことが分かった場合に、2つの列を1つの列 $b+c$ にすることを、「結合する」とする。
初期状態も、各頂点が「長さ1の列」と見なすことができる。

このとき、次のことが言える。

「列内の各要素 $i$ に対する $a_i$ の平均値」が最大の列は、列の最初の頂点の親を $p$ として、 $p$ が含まれる列の処理後、即座に選んで損しない。

略証
         k+1番目
① ...→ (c1 → ... → cm) → (b1 → ...... → bn) → ...
② ...→ (b1 → ...... → bn) → (c1 → ... → cm) → ...

列 b を、今ある中で「ai の平均値が最大の列」とする。
①において並び順で b の直前にある列を c とする。
c には b1 の親は含まれていない前提のもと、
①の状況は入れ替えて②にしても、並び順の制約には違反しない。

b の要素数を n、 ai の合計値を Sb とする。（つまり、平均値は Sb/n）
c の要素数を m、 ai の合計値を Sc とする。

この2列だけに着目したコストは、

  ① Σ_{i=1～m} (k+i) a[ci] + Σ_{i=1～n} (k+i+m) a[bi]
  ② Σ_{i=1～n} (k+i) a[bi] + Σ_{i=1～m} (k+i+n) a[ci]

これが ① >= ② であることを示す。
同じものを消せば
   m * Sb >=  n * Sc
つまり Sb/n >= Sc/m であればよいことになる。
これは b が平均値最大であるという前提から、常に成り立つ。

よって①のような状況では、②に置き換えて損しないことが言える。

これを繰り返すと、
「列 b は、b1 の親 p が含まれる列の直後に持ってきてよい」となり、
(..., p, ..., b1, b2, ..., bn) という1つの列に結合できる。

$N$ 回、列の結合を繰り返すことで、最適な $q$ が構築できる。
コストだけなら、差分に着目すれば具体的な $q$ を求めなくともよい。

$f(b)$ を「列 $b$ だけを1番目から操作した時のコスト $\sum_{i=1}^{n}ia_{bi}$ 」とする。
列を $b→c$ の順に結合するとき、 $f(b)+f(c)$ は既に求まっているとして、 $f(b+c)$ は $(bの要素数)\times(cの合計)$ だけ増加する。（上記の式参照）

この増分を結合のたびに加算していけば、 $N$ 回の結合で結果的に $f(q)$ となる。

実装例としては、以下のようにすればよい。

列を平均値が大きな方から取り出せる、途中要素の変更も可能な優先度キューを $H$ とする。
- はじめ、(平均値, 列先頭の頂点番号) = $(a_i,i)$ （ $i=1～N$ ）が入っている。
以下を $N$ 回繰り返す。
- $H$ から平均値最大の列 $b$ を取り出す。 $a_i$ の合計値 $S_b$ 、要素数 $C_b$ とする。
- $b_1$ の親頂点が含まれる列 $d$ を特定する。要素数 $C_d$ とする。
- $C_dS_b$ をコストに加算する。
- $d$ と $b$ を結合し、 $d$ の平均値、 $S_d,C_d$ をそれぞれ更新する。

$b$ から $d$ の特定には、Union-Findが使える。
Uniteする際、必ず「列の先頭の頂点がその連結成分のリーダーとなる」ように結合すれば、次、 $d$ から更にその親を見つけるのが多少楽になる。

Python3

有理数の比較を行うので、（小数でも大丈夫とは思うが、一応）誤差が出ないように自前のPriorityQueueを実装した。
もとい、ChatGPTに実装してもらった。さすが、枯れていて実装例が豊富にあるアルゴリズムは問題なく再現してくれる。

        
          
          
              
              from operator import lt
 
 
class PriorityQueue:
    def __init__(self, cmp=lt):
        self.heap = []
        self.cmp = cmp
 
    def push(self, item):
        # 項目と優先度のタプルを追加し、ヒープ条件を満たすよう調整
        self.heap.append(item)
        self._sift_up(len(self.heap) - 1)
 
    def pop(self):
        if len(self.heap) == 0:
            raise IndexError("pop from an empty priority queue")
 
        # ルート要素を取り出し、最後の要素をルートに移動してヒープを再調整
        self._swap(0, len(self.heap) - 1)
        item = self.heap.pop()  # 最後の要素がルートだったのでそれを返す
        self._sift_down(0)
        return item  # 優先度を無視して、元のアイテムのみ返す
 
    def peek(self):
        if len(self.heap) == 0:
            raise IndexError("peek from an empty priority queue")
        return self.heap[0]  # ルートのアイテムを返す
 
    def _sift_up(self, idx):
        if idx == 0:
            return
        # 追加した要素を適切な位置まで持ち上げる
        parent = (idx - 1) // 2
        if idx > 0 and self.cmp(self.heap[idx], self.heap[parent]):
            self._swap(idx, parent)
            self._sift_up(parent)
 
    def _sift_down(self, idx):
        # 根からスタートし、下方向に適切な位置に移動させる
        child = 2 * idx + 1
        if child < len(self.heap):
            # 右の子が存在し、右の子の方が小さい場合はそちらと比較
            if child + 1 < len(self.heap) and self.cmp(self.heap[child + 1], self.heap[child]):
                child += 1
            # 現在の要素が最小子よりも大きい場合は交換
            if self.cmp(self.heap[child], self.heap[idx]):
                self._swap(idx, child)
                self._sift_down(child)
 
    def _swap(self, i, j):
        # ヒープ中の2つの要素を交換
        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]
 
    def __len__(self):
        return len(self.heap)
 
 
class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n
 
    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x
 
    def find(self, x, y):
        return self.root(x) == self.root(y)
 
    def unite(self, x, y):
        """ [この問題のための改変] 必ずxの方を新たな代表とする """
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return False
        self.table[r1] += self.table[r2]
        self.table[r2] = r1
        return True
 
    def get_size(self, x):
        return -self.table[self.root(x)]
 
 
def solve(n, ppp, aaa):
    MOD = 998244353
    s = sum(aaa)
    ans = 0
    uft = UnionFind(n + 1)
 
    def cmp(a, b):
        return a[0] * b[1] > a[1] * b[0]
 
    hq = PriorityQueue(cmp)
    current_sum = [0]
    current_cnt = [1]
    for i, a in enumerate(aaa, start=1):
        hq.push((a, 1, i))
        current_sum.append(a)
        current_cnt.append(1)
 
    for _ in range(n):
        while True:
            a, b, v = hq.pop()
            if current_sum[v] == a and current_cnt[v] == b:
                break
 
        u = uft.root(ppp[v - 1])
        ans += current_cnt[u] * a
        ans %= MOD
        current_sum[u] += a
        current_cnt[u] += b
        uft.unite(u, v)
        if u != 0:
            hq.push((current_sum[u], current_cnt[u], u))
 
    ans *= pow(s, -1, MOD)
    ans %= MOD
 
    return ans
 
 
t = int(input())
for _ in range(t):
    n = int(input())
    ppp = list(map(int, input().split()))
    aaa = list(map(int, input().split()))
    ans = solve(n, ppp, aaa)
    print(ans)

            

        
      