Problem:

G: Tree

問題

N 頂点からなる木が与えられます。木の各頂点には 1 から N の番号が割り振られています。また、N−1 本の辺のうち、 i\ (= 1, 2, ..., N-1) 本目の辺は頂点 u_i と頂点 v_i とを結んでいます。

この木の K 個の空でない部分グラフの組であって、各部分グラフがいずれも連結で、どの二つの相異なる部分グラフも頂点を共有しないものの個数を求めるプログラムを作成してください。ただし答えがとても大きくなることがあるため、998244353 で割ったあまりを答えてください。

なお、K 個の部分グラフからなる集合が同一のものは、K 個の部分グラフの順序が異なるものも同一視するものとします。

入力形式

N K
u_1 v_1
:
u_{N-1} v_{N-1}

制約

 2 \leq N \leq 10^{5}

 1 \leq K \leq min ( N, 300 )

 1 \leq u_i, v_i \leq N

 u_i \neq v_i

 i \neq j に対して (u_i, v_i) \neq (u_j, v_j)

 入力はすべて整数で与えられる。

 与えられるグラフは木であることが保証される。

出力形式

答えを表す整数を一行に出力してください。998244353 で割ったあまりを出力することに注意してください。

入力例 1

3 2
1 2
1 3

出力例 1

5

以下の 5 通りがあります:

 \{ 1 \} と \{ 2 \}

 \{ 1 \} と \{ 3 \}

 \{ 2 \} と \{ 3 \}

 \{ 1, 2 \} と \{ 3 \}

 \{ 1, 3 \} と \{ 2 \}

入力例 2

4 4
1 2
1 3
1 4

出力例 2

1

( \{ 1 \}, \{ 2 \}, \{ 3 \}, \{ 4 \} ) の  1 通りしかありません。

入力例 3

7 4
1 7
2 1
7 4
3 4
5 7
6 3

出力例 3

166



Sample input:
3 2
1 2
1 3


Sample output:
5

Submission status: Accepted

Solution code:
import sys
        
def getpar(Edge, p):
    N = len(Edge)
    par = [0]*N
    par[0] = -1
    par[p]  -1
    stack = [p]
    visited = set([p])
    while stack:
        vn = stack.pop()
        for vf in Edge[vn]:
            if vf in visited:
                continue
            visited.add(vf)
            par[vf] = vn
            stack.append(vf)
    return par

def topological_sort_tree(E, r):
    Q = [r]
    L = []
    visited = set([r])
    while Q:
        vn = Q.pop()
        L.append(vn)
        for vf in E[vn]:
            if vf not in visited:
                visited.add(vf)
                Q.append(vf)
    return L

mod = 998244353
N, K = map(int, input().split())

seg = 110
lim = K+1
M = (1<<(lim*seg)) - 1
segb = (1<<seg)-1
fold = 47
tm = (1<<fold)%mod
bfilter = 0


for _ in range(lim):
    bfilter = bfilter<<seg
    bfilter = bfilter | ((1<<fold) -1)
cfilter = M ^ bfilter
def modulo(x):
    x = x&M
    b = x&bfilter
    c = ((x&cfilter)>>fold) * tm
    x = b+c
    b = x&bfilter
    c = ((x&cfilter)>>fold) * tm
    x = b+c
    b = x&bfilter
    c = ((x&cfilter)>>fold) * tm
    x = b+c
    return x
    
        
Edge = [[] for _ in range(N)]
D = [0]*N
for _ in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    a -= 1
    b -= 1
    Edge[a].append(b)
    Edge[b].append(a)

P = getpar(Edge, 0)
L = topological_sort_tree(Edge, 0)

dp1 = [(1<<seg)]*N
dp2 = [1]*N
for l in L[:0:-1]:
    p = P[l]
    dp1[p] = modulo(dp1[p]*((dp1[l]>>seg) + dp1[l] + dp2[l]))
    dp2[p] = modulo(dp2[p]*(dp1[l]+dp2[l]))
res1 = (dp1[0]>>(K*seg)) & segb
res2 = (dp2[0]>>(K*seg)) & segb
print((res1+res2)%mod)