重心分解

木の重心(取り除くと、残りの各部分の頂点数が全体の半分以下になる頂点)を取り除き、残った各部分でも同じことを繰り返します。各頂点は $O(\log N)$ 個の部分にしか含まれないので、各部分で「重心を通るパス」を部分の大きさに比例する時間で数えれば、全体で $O(N \log N)$ になります。再帰は使いません。

重心 $c$ を取り除くときの部分は、 $c$ から level が level[c] より大きい頂点だけをたどって行ける頂点全体です。

例(長さがちょうど $K$ のパスの本数):

cpar, level = centroid_decomposition(g)
ans = 0
for c in range(n):
    cnt = [1]  # cnt[d] = これまでに見た頂点のうち、c からの距離が d のものの個数
    for s in g[c]:
        if level[s] < level[c]:
            continue
        ds = []
        q = [(s, c, 1)]
        for v, p, d in q:
            ds.append(d)
            for to in g[v]:
                if to != p and level[to] > level[c]:
                    q.append((to, v, d + 1))
        for d in ds:
            if d <= K and K - d < len(cnt):
                ans += cnt[K - d]
        for d in ds:
            if d == len(cnt):
                cnt.append(0)
            cnt[d] += 1

コード

def centroid_decomposition(g):
    # 木の隣接リスト g から (cpar, level) を返す
    # cpar[v] = 重心分解の木での v の親(根は -1)、level[v] = その木での v の深さ
    n = len(g)
    cpar = [-1] * n
    level = [-1] * n  # -1 はまだ取り除いていない頂点
    bpar = [-1] * n
    size = [1] * n
    st = [(0, -1, 0)]  # (成分の頂点, 親になる重心, 深さ)
    while st:
        s, p, d = st.pop()
        bpar[s] = -1
        comp = [s]
        for v in comp:
            size[v] = 1
            for to in g[v]:
                if to != bpar[v] and level[to] == -1:
                    bpar[to] = v
                    comp.append(to)
        for v in reversed(comp[1:]):
            size[bpar[v]] += size[v]
        c = s
        while True:
            for to in g[c]:
                if to != bpar[c] and level[to] == -1 and 2 * size[to] > len(comp):
                    c = to
                    break
            else:
                break
        cpar[c] = p
        level[c] = d
        for to in g[c]:
            if level[to] == -1:
                st.append((to, c, d + 1))
    return cpar, level