Distance in Tree (Codeforces 161D)
Tags: DP, trees

A tree is a connected graph that doesn't contain any cycles.The distance between two vertices of a tree is the length (in edges) of the shortest path between these vertices.You are given a tree with n vertices and a positive number k. Find the number of distinct pairs of the vertices which have a distance of exactly k between them. Note that pairs (v, u) and (u, v) are considered to be the same pair.

Input: The first line contains two integers n and k (1 ≤ n ≤ 50000, 1 ≤ k ≤ 500) — the number of vertices and the required distance between the vertices.Next n - 1 lines describe the edges as "ai bi" (without the quotes) (1 ≤ ai, bi ≤ n, ai ≠ bi), where ai and bi are the vertices connected by the i-th edge. All given edges are different.

Output: Print a single integer — the number of distinct pairs of the tree's vertices which have a distance of exactly k between them.Please do not use the %lld specifier to read or write 64-bit integers in С++. It is preferred to use the cin, cout streams or the %I64d specifier.

Note: In the first sample the pairs of vertexes at distance 2 from each other are (1, 3), (1, 5), (3, 5) and (2, 4).

Sample Inputs:
5 2
1 2
2 3
3 4
2 5

5 3
1 2
2 3
3 4
4 5

Sample Outputs:
4

2

Solution (Python):
(n, k) = map(int, input().split())
(t, q) = ([[] for i in range(n + 1)], [1])
for j in range(n - 1):
    (a, b) = map(int, input().split())
    t[a].append(b)
    t[b].append(a)
for x in q:
    for y in t[x]:
        t[y].remove(x)
    q.extend(t[x])
q.reverse()
(a, s) = ({}, 0)
for x in q:
    a[x] = [1]
    u = len(a[x])
    for y in t[x]:
        v = len(a[y])
        for d in range(max(0, k - u), v):
            s += a[y][d] * a[x][k - d - 1]
        if v >= u:
            for d in range(u - 1):
                a[x][d + 1] += a[y][d]
            a[x] += a[y][u - 1:]
            u = v + 1
        else:
            for d in range(0, v):
                a[x][d + 1] += a[y][d]
    if u > k:
        a[x].pop()
print(s)

Solution (Python):
(n, k) = map(int, input().split())
g = [[] for i in range(n)]
for _ in range(n - 1):
    (a, b) = map(int, input().split())
    g[a - 1].append(b - 1)
    g[b - 1].append(a - 1)
q = [0]
for x in q:
    for y in g[x]:
        g[y].remove(x)
    q.extend(g[x])
q.reverse()
(d, cnt) = ({}, 0)
for x in q:
    d[x] = [1]
    l_x = len(d[x])
    for to in g[x]:
        l_to = len(d[to])
        for i in range(max(0, k - l_x), l_to):
            cnt += d[to][i] * d[x][k - i - 1]
        if l_to >= l_x:
            for i in range(l_x - 1):
                d[x][i + 1] += d[to][i]
            d[x] += d[to][l_x - 1:]
            l_x = l_to + 1
        else:
            for i in range(l_to):
                d[x][i + 1] += d[to][i]
    if l_x > k:
        d[x].pop()
print(cnt)

Solution (Python):
(n, k) = map(int, input().split())
(t, q) = ([[] for i in range(n + 1)], [1])
for j in range(n - 1):
    (a, b) = map(int, input().split())
    t[a].append(b)
    t[b].append(a)
for x in q:
    for y in t[x]:
        t[y].remove(x)
    q.extend(t[x])
q.reverse()
(a, luismore) = ({}, 0)
for x in q:
    a[x] = [1]
    u = len(a[x])
    for y in t[x]:
        v = len(a[y])
        for d in range(max(0, k - u), v):
            luismore += a[y][d] * a[x][k - d - 1]
        if v >= u:
            for d in range(u - 1):
                a[x][d + 1] += a[y][d]
            a[x] += a[y][u - 1:]
            u = v + 1
        else:
            for d in range(0, v):
                a[x][d + 1] += a[y][d]
    if u > k:
        a[x].pop()
print(luismore)