結果
| 問題 | No.3443 Sum of (Tree Distances)^K 1 |
| コンテスト | |
| ユーザー |
 tassei903
|
| 提出日時 | 2026-02-10 18:59:44 |
| 言語 | PyPy3 (7.3.17) |
| 結果 |
AC
|
| 実行時間 | 376 ms / 2,000 ms |
| コード長 | 6,792 bytes |
| 記録 | |
| コンパイル時間 | 545 ms |
| コンパイル使用メモリ | 82,088 KB |
| 実行使用メモリ | 124,060 KB |
| 最終ジャッジ日時 | 2026-02-10 19:00:00 |
| 合計ジャッジ時間 | 14,952 ms |
|
ジャッジサーバーID （参考情報） |
judge2 / judge3 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 3 |
| other | AC * 47 |
ソースコード
import sys
input = lambda :sys.stdin.readline()[:-1]
ni = lambda :int(input())
na = lambda :list(map(int,input().split()))
yes = lambda :print("yes");Yes = lambda :print("Yes")
no = lambda :print("no");No = lambda :print("No")
inf = 10**16
#######################################################################
mod = 998244353
nn = 2 * 10 ** 6 + 100
fact = [1] * nn
for i in range(nn - 1):
fact[i + 1] = fact[i] * (i + 1) % mod
invfact = [1] * nn
invfact[nn - 1] = pow(fact[nn - 1], mod - 2, mod)
for i in range(nn - 1)[::-1]:
invfact[i] = invfact[i + 1] * (i + 1) % mod
def binom(x, y):
if x < 0 or y < 0 or x - y < 0:
return 0
return fact[x] * invfact[y] % mod * invfact[x - y] % mod
def path(x):
if x == 1:
return 1
else:
return fact[x] * pow(2, mod-2, mod) % mod
MOD = 998244353
_IMAG = 911660635
_IIMAG = 86583718
_rate2 = (0, 911660635, 509520358, 369330050, 332049552, 983190778, 123842337, 238493703, 975955924, 603855026, 856644456, 131300601, 842657263, 730768835, 942482514, 806263778, 151565301, 510815449, 503497456, 743006876, 741047443, 56250497, 867605899, 0)
_irate2 = (0, 86583718, 372528824, 373294451, 645684063, 112220581, 692852209, 155456985, 797128860, 90816748, 860285882, 927414960, 354738543, 109331171, 293255632, 535113200, 308540755, 121186627, 608385704, 438932459, 359477183, 824071951, 103369235, 0)
_rate3 = (0, 372528824, 337190230, 454590761, 816400692, 578227951, 180142363, 83780245, 6597683, 70046822, 623238099, 183021267, 402682409, 631680428, 344509872, 689220186, 365017329, 774342554, 729444058, 102986190, 128751033, 395565204, 0)
_irate3 = (0, 509520358, 929031873, 170256584, 839780419, 282974284, 395914482, 444904435, 72135471, 638914820, 66769500, 771127074, 985925487, 262319669, 262341272, 625870173, 768022760, 859816005, 914661783, 430819711, 272774365, 530924681, 0)
def _fft(a):
n = len(a)
h = (n - 1).bit_length()
le = 0
for le in range(0, h - 1, 2):
p = 1 << (h - le - 2)
rot = 1
for s in range(1 << le):
rot2 = rot * rot % MOD
rot3 = rot2 * rot % MOD
offset = s << (h - le)
for i in range(p):
a0 = a[i + offset]
a1 = a[i + offset + p] * rot
a2 = a[i + offset + p * 2] * rot2
a3 = a[i + offset + p * 3] * rot3
a1na3imag = (a1 - a3) % MOD * _IMAG
a[i + offset] = (a0 + a2 + a1 + a3) % MOD
a[i + offset + p] = (a0 + a2 - a1 - a3) % MOD
a[i + offset + p * 2] = (a0 - a2 + a1na3imag) % MOD
a[i + offset + p * 3] = (a0 - a2 - a1na3imag) % MOD
rot = rot * _rate3[(~s & -~s).bit_length()] % MOD
if h - le & 1:
rot = 1
for s in range(1 << (h - 1)):
offset = s << 1
l = a[offset]
r = a[offset + 1] * rot
a[offset] = (l + r) % MOD
a[offset + 1] = (l - r) % MOD
rot = rot * _rate2[(~s & -~s).bit_length()] % MOD
def _ifft(a):
n = len(a)
h = (n - 1).bit_length()
le = h
for le in range(h, 1, -2):
p = 1 << (h - le)
irot = 1
for s in range(1 << (le - 2)):
irot2 = irot * irot % MOD
irot3 = irot2 * irot % MOD
offset = s << (h - le + 2)
for i in range(p):
a0 = a[i + offset]
a1 = a[i + offset + p]
a2 = a[i + offset + p * 2]
a3 = a[i + offset + p * 3]
a2na3iimag = (a2 - a3) * _IIMAG % MOD
a[i + offset] = (a0 + a1 + a2 + a3) % MOD
a[i + offset + p] = (a0 - a1 + a2na3iimag) * irot % MOD
a[i + offset + p * 2] = (a0 + a1 - a2 - a3) * irot2 % MOD
a[i + offset + p * 3] = (a0 - a1 - a2na3iimag) * irot3 % MOD
irot = irot * _irate3[(~s & -~s).bit_length()] % MOD
if le & 1:
p = 1 << (h - 1)
for i in range(p):
l = a[i]
r = a[i + p]
a[i] = l + r if l + r < MOD else l + r - MOD
a[i + p] = l - r if l - r >= 0 else l - r + MOD
def ntt(a) -> None:
if len(a) <= 1: return
_fft(a)
def intt(a) -> None:
if len(a) <= 1: return
_ifft(a)
iv = pow(len(a), MOD - 2, MOD)
for i, x in enumerate(a): a[i] = x * iv % MOD
def multiply(s: list, t: list) -> list:
n, m = len(s), len(t)
l = n + m - 1
if min(n, m) <= 60:
a = [0] * l
for i, x in enumerate(s):
for j, y in enumerate(t):
a[i + j] += x * y
return [x % MOD for x in a]
z = 1 << (l - 1).bit_length()
a = s + [0] * (z - n)
b = t + [0] * (z - m)
_fft(a)
_fft(b)
for i, x in enumerate(b): a[i] = a[i] * x % MOD
_ifft(a)
a[l:] = []
iz = pow(z, MOD - 2, MOD)
return [x * iz % MOD for x in a]
def pow2(a: list) -> list:
l = (len(a) << 1) - 1
if len(a) <= 60:
s = [0] * l
for i, x in enumerate(a):
for j, y in enumerate(a):
s[i + j] += x * y
return [x % MOD for x in s]
k = 2; M = 4
while M < l: M <<= 1; k += 1
s = a + [0] * (M - len(a))
_fft(s, k)
s = [x * x % MOD for x in s]
_ifft(s, k)
s[l:] = []
invm = pow(M, MOD - 2, MOD)
return [x * invm % MOD for x in s]
def ntt_doubling(a: list) -> None:
M = len(a)
intt(a)
r = 1
zeta = pow(3, (MOD - 1) // (M << 1), MOD)
for i, x in enumerate(a):
a[i] = x * r % MOD
r = r * zeta % MOD
ntt(a)
"""
[1, 1, 3]
a を固定する
sum_{T 木} sum_{p in T} (p が 全て a以下なら ) ^ K
p を固定する
k = |p| pを含むような木が何個あるか
k * n^{n-k-1}
全てa以下で，長さがkのパスの個数
c(a, k) * k!/2
a!/(a-k)! * k * n ^ {n - k - 1} * (k-1) ^ K / 2
a! * n^(n-1)/2 * (a-k)! * k * (k-1)^K * n^{-k}
"""
def naive(n, K):
ans = [0] * (n + 1)
for a in range(1, n + 1):
for k in range(2, a + 1):
ans[a] += binom(a, k) * path(k) * k * pow(n, n - k - 1, mod) * pow(k-1, K, mod)
ans[a] %= mod
return ans
def solve(n, K):
f = [0] * (n + 1)
g = [0] * (n + 1)
ninv = pow(n, mod-2, mod)
z = 1
for i in range(n + 1):
f[i] = invfact[i]
g[i] = i * pow(i - 1, K, mod) % mod * z % mod
z = ninv * z % mod
# print(f, g)
h = multiply(f, g)
# print(h)
Z = pow(n, n - 1, mod) * pow(2, mod-2, mod) % mod
for i in range(1, n + 1):
h[i] *= fact[i] * Z % mod
h[i] %= mod
for i in range(1, n + 1):
print((h[i] - h[i-1]) % mod)
n, k = na()
# print(naive(n, k))
solve(n, k)