結果

問題 No.3449 Mex of Subtree
コンテスト
ユーザー 👑 potato167
提出日時 2026-01-04 17:25:08
言語 PyPy3
(7.3.17)
結果
AC  
実行時間 487 ms / 2,000 ms
コード長 2,014 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 231 ms
コンパイル使用メモリ 82,324 KB
実行使用メモリ 79,124 KB
最終ジャッジ日時 2026-02-20 20:51:46
合計ジャッジ時間 9,136 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 59
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

import sys

MOD = 998244353

def main():
    sys.setrecursionlimit(1_000_000)
    N = int(sys.stdin.readline())
    P = []
    if N >= 2:
        P = list(map(int, sys.stdin.readline().split()))

    children = [[] for _ in range(N + 1)]
    for i in range(2, N + 1):
        parent = P[i - 2]
        children[parent].append(i)

    # subtree sizes (parents have smaller indices: process from N down to 1)
    sz = [1] * (N + 1)
    for v in range(N, 0, -1):
        s = 1
        for c in children[v]:
            s += sz[c]
        sz[v] = s

    # to reduce intermediate polynomial growth, multiply children in increasing subtree size
    for v in range(1, N + 1):
        children[v].sort(key=lambda x: sz[x])

    dp = [None] * (N + 1)

    # convolution truncated to limit
    def convolve_trunc(a, b, limit):
        # commutative: iterate smaller outer to reduce Python overhead
        if len(a) > len(b):
            a, b = b, a
        max_len = min(limit, len(a) + len(b) - 2) + 1
        res = [0] * max_len
        b_len = len(b)
        for i, ai in enumerate(a):
            if ai == 0:
                continue
            maxj = min(b_len - 1, limit - i)
            # j from 0..maxj
            for j in range(maxj + 1):
                bj = b[j]
                if bj:
                    res[i + j] = (res[i + j] + ai * bj) % MOD
        return res

    # bottom-up DP
    for v in range(N, 0, -1):
        limit = sz[v]
        conv = [1]  # polynomial for sum of children's S
        for c in children[v]:
            conv = convolve_trunc(conv, dp[c], limit)
        # dp_v[s] = prefix sum of conv up to s
        dpv = [0] * (limit + 1)
        run = 0
        for s in range(limit + 1):
            if s < len(conv):
                run += conv[s]
            dpv[s] = run % MOD
        dp[v] = dpv

        # free children dp (tree: used only once)
        for c in children[v]:
            dp[c] = None

    ans = sum(dp[1]) % MOD
    print(ans)

if __name__ == "__main__":
    main()
0