UNICORNプログラミングコンテスト2022(AtCoder Beginner Contest 269) G,Ex問題メモ

import os
import sys

import numpy as np


def solve(inp):
    n, m = inp[:2]
    aaa = inp[2::2]
    bbb = inp[3::2]

    s = aaa.sum()
    diff = bbb - aaa

    diff_count = {n: n}
    diff_count.clear()
    for d in diff:
        if d in diff_count:
            diff_count[d] += 1
        else:
            diff_count[d] = 1

    INF = 1 << 60
    dp = np.full(m + 1, INF, np.int64)
    dp[s] = 0

    for d, x in diff_count.items():
        if d == 0:
            continue
        k = 1
        while x:
            ndp = dp.copy()
            c = min(k, x)
            if d > 0:
                for a in range(m + 1 - d * c):
                    b = a + d * c
                    ndp[b] = min(ndp[b], dp[a] + c)
            else:
                for a in range(-d * c, m + 1):
                    b = a + d * c
                    ndp[b] = min(ndp[b], dp[a] + c)
            dp = ndp
            x -= c
            k <<= 1

    dp[dp == INF] = -1

    return dp


SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print('\n'.join(map(str, ans)))

from heapq import heapify, heappop, heappush
from typing import List

import numpy as np


def heavy_light_decomposition(n: int, children: List[List[int]], root: int = 0) -> List[List[int]]:
    """
    ①--②--④--⑦--⑨  →  [[1,2,4,7,9], [10], [3,5,8], [6]]
     `--③--⑤-⑧ `-⑩
         `--⑥
    Heavy-path のリストに分解する。
    Heavy-path の並びは、（最も重い子を最初に訪れるような）オイラーツアーの訪問順となる。
    """
    weights = [-1] * n
    q = [root]
    while q:
        u = q[-1]
        if weights[u] == -1:
            weights[u] = -2
            q.extend(children[u])
        else:
            q.pop()
            weights[u] = 1 + sum(weights[v] for v in children[u])

    q = [root]
    progress = [0] * n
    current_path = []
    result = [current_path]
    while q:
        u = q[-1]
        if progress[u] == 0:
            children[u].sort(key=weights.__getitem__, reverse=True)
            current_path.append(u)
        if progress[u] >= len(children[u]):
            q.pop()
            continue
        v = children[u][progress[u]]
        if progress[u] > 0:
            current_path = []
            result.append(current_path)
        progress[u] += 1
        q.append(v)

    return result


def _convolve(f, g):
    # 小さいうちは愚直の方が高速。convolve_mod()と併用するなら大丈夫だが、単独で使用する場合はオーバーフローに注意
    if len(f) * len(g) < 10000:
        return np.convolve(f, g)

    fft_len = 1
    true_len = len(f) + len(g) - 1
    while fft_len < true_len:
        fft_len <<= 1

    Ff = np.fft.rfft(f, fft_len)
    Fg = np.fft.rfft(g, fft_len)
    Fh = Ff * Fg

    h = np.fft.irfft(Fh, fft_len)
    h = np.rint(h).astype(np.int64)

    return h[:true_len]


def convolve_mod(f, g, p):
    f1, f2 = np.divmod(f, 1 << 15)
    g1, g2 = np.divmod(g, 1 << 15)

    a = _convolve(f1, g1) % p
    c = _convolve(f2, g2) % p
    b = (_convolve(f1 + f2, g1 + g2) - (a + c)) % p

    h = (a << 30) + (b << 15) + c
    return h % p


n = int(input())
ppp = list(map(int, input().split()))
children = [[] for _ in range(n)]
parent = [-1] * n
for i, p in enumerate(ppp, start=1):
    p -= 1
    children[p].append(i)
    parent[i] = p
MOD = 998244353

hld_paths = heavy_light_decomposition(n, children, 0)

light_children_results = [[] for _ in range(n)]
default_dp1 = np.array([1], np.int64)
default_dp2 = np.array([0, 1], np.int64)
ans = []

for heavy_path in reversed(hld_paths):
    merge_array = []
    for v in heavy_path:
        if len(light_children_results[v]) == 0:
            merge_array.append((default_dp1.copy(), default_dp2.copy()))
        elif len(light_children_results[v]) == 1:
            lcr = light_children_results[v][0]
            merge_array.append((lcr, default_dp2.copy()))
        else:
            merge2_array = light_children_results[v]
            merge2_queue = []
            for i, lcr in enumerate(merge2_array):
                merge2_queue.append((len(lcr), i))
            heapify(merge2_queue)
            while len(merge2_queue) > 1:
                l1, i1 = heappop(merge2_queue)
                l2, i2 = heappop(merge2_queue)
                lcr3 = convolve_mod(merge2_array[i1], merge2_array[i2], MOD)
                i3 = len(merge2_array)
                merge2_array.append(lcr3)
                heappush(merge2_queue, (len(lcr3), i3))
            lcr = merge2_array[-1]
            merge_array.append((lcr, default_dp2.copy()))

    while len(merge_array) > 1:
        new_array = []
        l = len(merge_array)
        for i in range(0, l, 2):
            if i + 1 == l:
                new_array.append(merge_array[i])
                continue
            dp11, dp12 = merge_array[i]
            dp21, dp22 = merge_array[i + 1]

            dp31 = convolve_mod(dp11, dp21, MOD)
            dp32 = convolve_mod(dp11, dp22, MOD)
            dp32[:len(dp12)] += dp12
            dp32 %= MOD
            new_array.append((dp31, dp32))
        merge_array = new_array

    res1, res2 = merge_array[0]
    if len(res1) >= len(res2):
        res1[:len(res2)] += res2
        res = res1 % MOD
    else:
        res2[:len(res1)] += res1
        res = res2 % MOD

    leader = heavy_path[0]
    if leader != 0:
        p = parent[leader]
        light_children_results[p].append(res)

    else:
        ans = res

ans = ans.tolist()
if len(ans) < n + 1:
    ans.extend([0] * (n + 1 - len(ans)))

print('\n'.join(map(str, ans[1:])))