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

有理数の比較を行うので、（小数でも大丈夫とは思うが、一応）誤差が出ないように自前の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)