Equidistant Vertices (Codeforces 1551F)
Tags: DP, trees

A tree is an undirected connected graph without cycles.You are given a tree of $$$n$$$ vertices. Find the number of ways to choose exactly $$$k$$$ vertices in this tree (i. e. a $$$k$$$-element subset of vertices) so that all pairwise distances between the selected vertices are equal (in other words, there exists an integer $$$c$$$ such that for all $$$u, v$$$ ($$$u \ne v$$$, $$$u, v$$$ are in selected vertices) $$$d_{u,v}=c$$$, where $$$d_{u,v}$$$ is the distance from $$$u$$$ to $$$v$$$).Since the answer may be very large, you need to output it modulo $$$10^9 + 7$$$.

Input: The first line contains one integer $$$t$$$ ($$$1 \le t \le 10$$$) — the number of test cases. Then $$$t$$$ test cases follow.Each test case is preceded by an empty line.Each test case consists of several lines. The first line of the test case contains two integers $$$n$$$ and $$$k$$$ ($$$2 \le k \le n \le 100$$$) — the number of vertices in the tree and the number of vertices to be selected, respectively. Then $$$n - 1$$$ lines follow, each of them contains two integers $$$u$$$ and $$$v$$$ ($$$1 \le u, v \le n$$$, $$$u \neq v$$$) which describe a pair of vertices connected by an edge. It is guaranteed that the given graph is a tree and has no loops or multiple edges.

Output: For each test case output in a separate line a single integer — the number of ways to select exactly $$$k$$$ vertices so that for all pairs of selected vertices the distances between the vertices in the pairs are equal, modulo $$$10^9 + 7$$$ (in other words, print the remainder when divided by $$$1000000007$$$).


Sample Inputs:
3

4 2
1 2
2 3
2 4

3 3
1 2
2 3

5 3
1 2
2 3
2 4
4 5

Sample Outputs:
6
0
1

Solution (Python):
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
from collections import deque
M = 10 ** 9 + 7

def gettot(dic, k):
    nums = []
    for key in dic:
        nums.append(dic[key])
    m = len(nums)
    dp = [[0 for j in range(k + 1)] for i in range(m + 1)]
    for i in range(m + 1):
        dp[i][0] = 1
    for i in range(m - 1, -1, -1):
        for j in range(k + 1):
            if m - i < j:
                break
            dp[i][j] = dp[i + 1][j]
            if j > 0:
                dp[i][j] += nums[i] * dp[i + 1][j - 1] % M
            dp[i][j] = dp[i][j] % M
    return dp[0][k]

def main(t):
    s = input()
    (n, k) = map(int, input().split())
    neigh = [[] for i in range(n)]
    for _ in range(n - 1):
        (u, v) = map(int, input().split())
        neigh[u - 1].append(v - 1)
        neigh[v - 1].append(u - 1)
    ans = 0
    if k == 2:
        print(n * (n - 1) // 2)
        return
    for root in range(n):
        queue = deque()
        layers = [{}]
        for ele in neigh[root]:
            queue.append([ele, ele, root, 0])
            layers[0][ele] = 1
        while queue:
            [index, ancestor, father, l] = queue.popleft()
            for nextindex in neigh[index]:
                if nextindex == father:
                    continue
                if l + 1 == len(layers):
                    layers.append({})
                if ancestor not in layers[l + 1]:
                    layers[l + 1][ancestor] = 0
                layers[l + 1][ancestor] += 1
                queue.append([nextindex, ancestor, index, l + 1])
        for dic in layers:
            if len(dic) < k:
                break
            ans += gettot(dic, k)
            ans = ans % M
    print(ans)
T = int(input())
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
from collections import deque
M = 10 ** 9 + 7

def gettot(dic, k):
    nums = []
    for key in dic:
        nums.append(dic[key])
    m = len(nums)
    dp = [[0 for j in range(k + 1)] for i in range(m + 1)]
    for i in range(m + 1):
        dp[i][0] = 1
    for i in range(m - 1, -1, -1):
        for j in range(k + 1):
            if m - i < j:
                break
            dp[i][j] = dp[i + 1][j]
            if j > 0:
                dp[i][j] += nums[i] * dp[i + 1][j - 1] % M
            dp[i][j] = dp[i][j] % M
    return dp[0][k]

def main(t):
    s = input()
    (n, k) = map(int, input().split())
    neigh = [[] for i in range(n)]
    for _ in range(n - 1):
        (u, v) = map(int, input().split())
        neigh[u - 1].append(v - 1)
        neigh[v - 1].append(u - 1)
    ans = 0
    if k == 2:
        print(n * (n - 1) // 2)
        return
    for root in range(n):
        queue = deque()
        layers = [{}]
        for ele in neigh[root]:
            queue.append([ele, ele, root, 0])
            layers[0][ele] = 1
        while queue:
            [index, ancestor, father, l] = queue.popleft()
            for nextindex in neigh[index]:
                if nextindex == father:
                    continue
                if l + 1 == len(layers):
                    layers.append({})
                if ancestor not in layers[l + 1]:
                    layers[l + 1][ancestor] = 0
                layers[l + 1][ancestor] += 1
                queue.append([nextindex, ancestor, index, l + 1])
        for dic in layers:
            if len(dic) < k:
                break
            ans += gettot(dic, k)
            ans = ans % M
    print(ans)
T = int(input())
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
from collections import deque
M = 10 ** 9 + 7

def gettot(dic, k):
    nums = []
    for key in dic:
        nums.append(dic[key])
    m = len(nums)
    dp = [[0 for j in range(k + 1)] for i in range(m + 1)]
    for i in range(m + 1):
        dp[i][0] = 1
    for i in range(m - 1, -1, -1):
        for j in range(k + 1):
            if m - i < j:
                break
            dp[i][j] = dp[i + 1][j]
            if j > 0:
                dp[i][j] += nums[i] * dp[i + 1][j - 1] % M
            dp[i][j] = dp[i][j] % M
    return dp[0][k]

def main(t):
    s = input()
    (n, k) = map(int, input().split())
    neigh = [[] for i in range(n)]
    for _ in range(n - 1):
        (u, v) = map(int, input().split())
        neigh[u - 1].append(v - 1)
        neigh[v - 1].append(u - 1)
    ans = 0
    if k == 2:
        print(n * (n - 1) // 2)
        return
    for root in range(n):
        queue = deque()
        layers = [{}]
        for ele in neigh[root]:
            queue.append([ele, ele, root, 0])
            layers[0][ele] = 1
        while queue:
            [index, ancestor, father, l] = queue.popleft()
            for nextindex in neigh[index]:
                if nextindex == father:
                    continue
                if l + 1 == len(layers):
                    layers.append({})
                if ancestor not in layers[l + 1]:
                    layers[l + 1][ancestor] = 0
                layers[l + 1][ancestor] += 1
                queue.append([nextindex, ancestor, index, l + 1])
        for dic in layers:
            if len(dic) < k:
                break
            ans += gettot(dic, k)
            ans = ans % M
    print(ans)
T = int(input())
t = 1
while t <= T:
    main(t)
    t += 1