Apple Tree (Codeforces 1843D)
Tags: Math, trees

Timofey has an apple tree growing in his garden; it is a rooted tree of $$$n$$$ vertices with the root in vertex $$$1$$$ (the vertices are numbered from $$$1$$$ to $$$n$$$). A tree is a connected graph without loops and multiple edges.This tree is very unusual — it grows with its root upwards. However, it's quite normal for programmer's trees.The apple tree is quite young, so only two apples will grow on it. Apples will grow in certain vertices (these vertices may be the same). After the apples grow, Timofey starts shaking the apple tree until the apples fall. Each time Timofey shakes the apple tree, the following happens to each of the apples:Let the apple now be at vertex $$$u$$$.  If a vertex $$$u$$$ has a child, the apple moves to it (if there are several such vertices, the apple can move to any of them).  Otherwise, the apple falls from the tree. It can be shown that after a finite time, both apples will fall from the tree.Timofey has $$$q$$$ assumptions in which vertices apples can grow. He assumes that apples can grow in vertices $$$x$$$ and $$$y$$$, and wants to know the number of pairs of vertices ($$$a$$$, $$$b$$$) from which apples can fall from the tree, where $$$a$$$ — the vertex from which an apple from vertex $$$x$$$ will fall, $$$b$$$ — the vertex from which an apple from vertex $$$y$$$ will fall. Help him do this.

Input: The first line contains integer $$$t$$$ ($$$1 \leq t \leq 10^4$$$) — the number of test cases.The first line of each test case contains integer $$$n$$$ ($$$2 \leq n \leq 2 \cdot 10^5$$$) — the number of vertices in the tree.Then there are $$$n - 1$$$ lines describing the tree. In line $$$i$$$ there are two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i \ne v_i$$$) — edge in tree.The next line contains a single integer $$$q$$$ ($$$1 \leq q \leq 2 \cdot 10^5$$$) — the number of Timofey's assumptions.Each of the next $$$q$$$ lines contains two integers $$$x_i$$$ and $$$y_i$$$ ($$$1 \leq x_i, y_i \leq n$$$) — the supposed vertices on which the apples will grow for the assumption $$$i$$$.It is guaranteed that the sum of $$$n$$$ does not exceed $$$2 \cdot 10^5$$$. Similarly, It is guaranteed that the sum of $$$q$$$ does not exceed $$$2 \cdot 10^5$$$.

Output: For each Timofey's assumption output the number of ordered pairs of vertices from which apples can fall from the tree if the assumption is true on a separate line.

Note: In the first example:   For the first assumption, there are two possible pairs of vertices from which apples can fall from the tree: $$$(4, 4), (5, 4)$$$.  For the second assumption there are also two pairs: $$$(5, 4), (5, 5)$$$.  For the third assumption there is only one pair: $$$(4, 4)$$$.  For the fourth assumption, there are $$$4$$$ pairs: $$$(4, 4), (4, 5), (5, 4), (5, 5)$$$.   Tree from the first example. For the second example, there are $$$4$$$ of possible pairs of vertices from which apples can fall: $$$(2, 3), (2, 2), (3, 2), (3, 3)$$$. For the second assumption, there is only one possible pair: $$$(2, 3)$$$. For the third assumption, there are two pairs: $$$(3, 2), (3, 3)$$$.

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

2
5
5 1
1 2
2 3
4 3
2
5 5
5 1
5
3 2
5 3
2 1
4 2
3
4 3
2 1
4 2

Sample Outputs:
2
2
1
4
4
1
2

1
2
1
4
2

Solution (Python):
import sys
import threading
import collections
import datetime
stdin = sys.stdin
ni = lambda : int(ns())
na = lambda : list(map(int, stdin.readline().split()))
ns = lambda : stdin.readline()

def main():
    test = ni()
    for _ in range(test):
        n = ni()
        g = [[] for _ in range(n + 1)]
        for _ in range(n - 1):
            (x, y) = na()
            g[x].append(y)
            g[y].append(x)
        dp = [0 for _ in range(n + 1)]

        def dfs(node):
            stack = [node]
            visited = [False for _ in range(n + 1)]
            visited[node] = True
            while stack:
                x = stack[-1]
                discovered = True
                for y in g[x]:
                    if not visited[y]:
                        discovered = False
                        visited[y] = True
                        stack.append(y)
                if discovered:
                    leaves = 0
                    for y in g[x]:
                        leaves += dp[y]
                    dp[x] = max(1, leaves)
                    stack.pop()
        dfs(1)
        queries = ni()
        for _ in range(queries):
            (a, b) = na()
            print(dp[a] * dp[b])
if __name__ == '__main__':
    threading.stack_size(1 << 27)
    sys.setrecursionlimit(1 << 30)
    main_thread = threading.Thread(target=main)
    main_thread.start()
    main_thread.join()

Solution (Python):
import sys
import threading
sys.setrecursionlimit(10 ** 8)
threading.stack_size(2 * 10 ** 8)
input = sys.stdin.readline

def main():
    for _ in range(int(input())):
        n = int(input())
        tree = [[] for _ in range(n + 1)]
        for _ in range(n - 1):
            (x, y) = map(int, input().split())
            tree[x].append(y)
            tree[y].append(x)
        leaf_count = [0] * (n + 1)

        def dfs(node, parent):
            if node != 1 and len(tree[node]) == 1:
                leaf_count[node] = 1
            for child in tree[node]:
                if child != parent:
                    dfs(child, node)
                    leaf_count[node] += leaf_count[child]
        dfs(1, 0)
        for _ in range(int(input())):
            (x, y) = map(int, input().split())
            print(leaf_count[x] * leaf_count[y])
thread = threading.Thread(target=main)
thread.start()
thread.join()

Solution (Python):
import sys, math
import threading
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 8)

def main():

    def dfs(vertex, visited, dp):
        if visited[vertex] == 1:
            return 0
        elif vertex in dp:
            return dp[vertex]
        else:
            visited[vertex] = 1
            ans = 0
            if len(gp[vertex]) == 1 and vertex != 0:
                ans = 1
            for i in range(len(gp[vertex])):
                if visited[gp[vertex][i]] == 1:
                    continue
                else:
                    ans += dfs(gp[vertex][i], visited, dp)
            dp[vertex] = ans
            return ans
    for _ in range(int(input())):
        n = int(input())
        gp = []
        for i in range(n):
            gp.append([])
        for i in range(n - 1):
            (x, y) = map(int, input().split())
            gp[x - 1].append(y - 1)
            gp[y - 1].append(x - 1)
        visited = [0] * n
        dp = {}
        for i in range(int(input())):
            if i == 0:
                dfs(i, visited, dp)
            (x, y) = map(int, input().split())
            print(dp[x - 1] * dp[y - 1])
threading.stack_size(2 * 10 ** 8)
t = threading.Thread(target=main)
t.start()
t.join()