LuoTianyi and the Floating Islands (Easy Version) (Codeforces 1825D1)
Tags: Math, trees

This is the easy version of the problem. The only difference is that in this version $$$k\le\min(n,3)$$$. You can make hacks only if both versions of the problem are solved.  Chtholly and the floating islands. LuoTianyi now lives in a world with $$$n$$$ floating islands. The floating islands are connected by $$$n-1$$$ undirected air routes, and any two of them can reach each other by passing the routes. That means, the $$$n$$$ floating islands form a tree.One day, LuoTianyi wants to meet her friends: Chtholly, Nephren, William, .... Totally, she wants to meet $$$k$$$ people. She doesn't know the exact positions of them, but she knows that they are in pairwise distinct islands. She define an island is good if and only if the sum of the distances$$$^{\dagger}$$$ from it to the islands with $$$k$$$ people is the minimal among all the $$$n$$$ islands.Now, LuoTianyi wants to know that, if the $$$k$$$ people are randomly set in $$$k$$$ distinct of the $$$n$$$ islands, then what is the expect number of the good islands? You just need to tell her the expect number modulo $$$10^9+7$$$.$$$^{\dagger}$$$The distance between two islands is the minimum number of air routes you need to take to get from one island to the other.

Input: The first line contains two integers $$$n$$$ and $$$k$$$ ($$$1\le k \le \min(n,3), 1\le n \le 2\cdot 10^5$$$) — the number of the islands and people respectively.Next $$$n−1$$$ lines describe the air routes. The $$$i$$$-th of them contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \le u_i,v_i \le n, u_i \neq v_i$$$) — the islands connected by the $$$i$$$-th air route.

Output: Print a single integer — the expect number of the good islands modulo $$$10^9 + 7$$$.Formally, let $$$M = 10^9 + 7$$$. It can be shown that the answer can be expressed as an irreducible fraction $$$\frac{p}{q}$$$, where $$$p$$$ and $$$q$$$ are integers and $$$q \not \equiv 0$$$ ($$$\operatorname{mod} M$$$). Output the integer equal to $$$p \cdot q^{-1}$$$ $$$\operatorname{mod} M$$$. In other words, output such an integer $$$x$$$ that $$$0 \le x < M$$$ and $$$x \cdot q \equiv p$$$ ($$$\operatorname{mod} M$$$).

Note: In the first example the air routes form the following tree:   If the people are in the islands $$$1$$$ and $$$2$$$, then islands $$$1$$$ and $$$2$$$ will be good.The sum of the distances from island $$$1$$$ or $$$2$$$ to all the people is $$$1+0=1$$$, which is the minimal. While the sum of the distances from island $$$3$$$ to all the people is $$$2+1=3$$$, which is greater than $$$1$$$.Like this, when the people are in island $$$1$$$ and $$$3$$$, then islands $$$1,2$$$ and $$$3$$$ will be good.When the people are in islands $$$1$$$ and $$$4$$$, then islands $$$1,2,3$$$ and $$$4$$$ will be good.When the people are in islands $$$2$$$ and $$$3$$$, then islands $$$2$$$ and $$$3$$$ will be good.When the people are in islands $$$2$$$ and $$$4$$$, then islands $$$2,3$$$ and $$$4$$$ will be good.When the people are in islands $$$3$$$ and $$$4$$$, then islands $$$3$$$ and $$$4$$$ will be good.So the expect of the number of the good islands is $$$\frac{16}{6}$$$, which equals to $$$666666674$$$ modulo $$$10^9+7$$$.In the second example the air routes form the following tree:  There is always the only good island, so the expected number is $$$1$$$.

Sample Inputs:
4 2
1 2
2 3
3 4

5 1
1 2
2 3
3 4
3 5

Sample Outputs:
666666674

1

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()
from collections import deque

def solve():
    (n, k) = readInts()
    graph = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        (u, v) = readInts()
        graph[u].append(v)
        graph[v].append(u)
    if k == 1 or k == 3:
        return 1
    m = 10 ** 9 + 7
    st = [(1, 0)]
    par = {1: -1}
    vis = [False] * (n + 1)
    size = [1] * (n + 1)
    d = [0] * (n + 1)
    while st:
        (node, dis) = st[-1]
        if vis[node]:
            (v, _) = st.pop()
            if node == 1:
                continue
            size[par[v]] += size[node]
            d[1] += dis
            d[1] %= m
            continue
        vis[node] = True
        for nei in graph[node]:
            if nei in par:
                continue
            par[nei] = node
            st.append((nei, dis + 1))
    (a, b) = (n * (n - 1) + d[1], pow(n * (n - 1), m - 2, m))
    st = [1]
    vis = [False] * (n + 1)
    vis[1] = True
    while st:
        node = st.pop()
        if node > 1:
            d[node] = d[par[node]] + n - 2 * size[node]
            d[node] %= m
            a += d[node]
        for nei in graph[node]:
            if not vis[nei]:
                st.append(nei)
                vis[nei] = True
    return a % m * b % m
print(solve())

Solution (Python):
import collections
MOD = int(1000000000.0 + 7)

def fast_power(n, k):
    result = 1 % MOD
    while k != 0:
        if k & 1:
            result = result * n % MOD
        n = n * n % MOD
        k >>= 1
    return result

def get_subtree_size(graph, root):
    tree_size = collections.defaultdict(int)
    father = {root: None}
    queue = collections.deque([root])
    order = collections.deque([root])
    while queue:
        node = queue.popleft()
        for neighbor in graph[node]:
            if neighbor == father[node]:
                continue
            father[neighbor] = node
            queue.append(neighbor)
            order.appendleft(neighbor)
    for node in order:
        tree_size[node] += 1
        if father[node]:
            tree_size[father[node]] += tree_size[node]
    return tree_size

def main():
    (n, k) = map(int, input().split())
    if k != 2:
        return print(1)
    graph = collections.defaultdict(list)
    for _ in range(n - 1):
        (u, v) = map(int, input().split())
        graph[u].append(v)
        graph[v].append(u)
    tree_size = get_subtree_size(graph=graph, root=1)
    p = 0
    for node in range(1, n + 1):
        prefix_sum = n - tree_size[node]
        for neighbor in graph[node]:
            if tree_size[neighbor] > tree_size[node]:
                continue
            p = (p + tree_size[neighbor] * prefix_sum) % MOD
            prefix_sum += tree_size[neighbor]
        p = (p + prefix_sum) % MOD
    q = (n * (n - 1) >> 1) % MOD
    print(p * fast_power(q, MOD - 2) % MOD)
if __name__ == '__main__':
    main()

Solution (Python):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(500000)
from collections import deque
(n, k) = map(int, input().split())
mod = 10 ** 9 + 7
adj = [[] for _ in range(n)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    adj[u - 1].append(v - 1)
    adj[v - 1].append(u - 1)
if k == 1 or k == 3:
    print(1)
    exit()
depth = [0] * n
seen = [0] * n
size = [1] * n
pa = [-1] * n
que = deque()
que.append((0, 1))
seen[0] = 1
while que:
    (v, task) = que.pop()
    if task == 0:
        for w in adj[v]:
            if w != pa[v]:
                size[v] += size[w]
    else:
        que.append((v, 0))
        seen[v] = 1
        for w in adj[v]:
            if seen[w] == 0:
                que.append((w, 1))
                pa[w] = v
                depth[w] = depth[v] + 1
now = sum(depth) + n - 1
tot = now
val = [0] * n
val[0] = now
seen = [0] * n
que = deque()
que.append((0, 1))
seen[0] = 1
while que:
    (v, task) = que.pop()
    if task == 0:
        continue
    else:
        que.append((v, 0))
        seen[v] = 1
        for w in adj[v]:
            if seen[w] == 0:
                que.append((w, 1))
                val[w] = val[v] + (n - 2 * size[w])
                tot += val[w]
inv = pow(n * (n - 1), mod - 2, mod)
print(tot * inv % mod)