Count Paths (Codeforces 1923E)
Tags: DP, dsu

You are given a tree, consisting of $$$n$$$ vertices, numbered from $$$1$$$ to $$$n$$$. Every vertex is colored in some color, denoted by an integer from $$$1$$$ to $$$n$$$.A simple path of the tree is called beautiful if:   it consists of at least $$$2$$$ vertices;  the first and the last vertices of the path have the same color;  no other vertex on the path has the same color as the first vertex. Count the number of the beautiful simple paths of the tree. Note that paths are considered undirected (i. e. the path from $$$x$$$ to $$$y$$$ is the same as the path from $$$y$$$ to $$$x$$$).

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of testcases.The first line of each testcase contains a single integer $$$n$$$ ($$$2 \le n \le 2 \cdot 10^5$$$) — the number of vertices in the tree.The second line contains $$$n$$$ integers $$$c_1, c_2, \dots, c_n$$$ ($$$1 \le c_i \le n$$$) — the color of each vertex.The $$$i$$$-th of the next $$$n - 1$$$ lines contains two integers $$$v_i$$$ and $$$u_i$$$ ($$$1 \le v_i, u_i \le n$$$; $$$v_i \neq u_i$$$) — the $$$i$$$-th edge of the tree.The given edges form a valid tree. The sum of $$$n$$$ over all testcases doesn't exceed $$$2 \cdot 10^5$$$.

Output: For each testcase, print a single integer — the number of the beautiful simple paths of the tree.


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

Sample Outputs:
1
3
0
3

Solution (Python):
for _ in range(int(input())):
    n = int(input())
    c = list(map(int, input().split()))
    for i in range(n):
        c[i] -= 1
    tree = [[] for i in range(n)]
    for i in range(n - 1):
        (v, u) = map(int, input().split())
        tree[v - 1].append(u - 1)
        tree[u - 1].append(v - 1)
    mas = [0] * n
    ans = 0
    kol = [0] * n
    stack = [[0, 1]]
    paren = [-1] * n
    while True:
        if len(stack) == 0:
            break
        (el, pos) = stack.pop()
        if pos == 1:
            kol[el] = mas[c[el]]
            mas[c[el]] = 0
            stack.append([el, 2])
            for x in tree[el]:
                if x != paren[el]:
                    stack.append([el, 3])
                    stack.append([x, 1])
                    paren[x] = el
        elif pos == 3:
            ans += mas[c[el]]
            mas[c[el]] = 0
        else:
            mas[c[el]] = kol[el] + 1
    print(ans + sum(kol))

Solution (Python):
import sys
from collections import defaultdict
import threading
sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline

def dfs(n, p):
    global dp, ans
    smallest = -1
    for i in adj[n]:
        if i == p:
            continue
        dfs(i, n)
        if smallest == -1 or len(dp[smallest]) < len(dp[i]):
            smallest = i
    for i in adj[n]:
        if i != p and i != smallest:
            for (x, y) in dp[i].items():
                if x != c[n - 1]:
                    ans += dp[smallest].get(x, 0) * y
                dp[smallest][x] = dp[smallest].get(x, 0) + y
    if smallest != -1:
        (dp[smallest], dp[n]) = (dp[n], dp[smallest])
    ans += dp[n].get(c[n - 1], 0)
    dp[n][c[n - 1]] = 1

def main():
    global n, c, adj, dp, ans
    for _ in range(int(input())):
        n = int(input())
        c = list(map(int, input().split()))
        adj = [[] for _ in range(n + 1)]
        dp = [defaultdict() for _ in range(n + 1)]
        for _ in range(n - 1):
            (u, v) = map(int, input().split())
            adj[v].append(u)
            adj[u].append(v)
        ans = 0
        dfs(1, 0)
        print(ans)
t1 = threading.Thread(target=main)
threading.stack_size(200000000)
t1.start()

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

def iin():
    return int(ipt())

def lmin():
    return list(map(int, ipt().split()))
M = 10 ** 6
MOD = 998244353
for _ in range(iin()):
    N = iin()
    A = lmin()
    to = [0] * (2 * N + 5)
    ne = [0] * (2 * N + 5)
    he = [0] * (2 * N + 5)
    cnt = 0
    for _ in range(N - 1):
        (u, v) = lmin()
        (u, v) = (u - 1, v - 1)
        cnt += 1
        to[cnt] = v
        ne[cnt] = he[u]
        he[u] = cnt
        cnt += 1
        to[cnt] = u
        ne[cnt] = he[v]
        he[v] = cnt
    val = [0] * (N + 1)
    pval = [0] * (N + 1)
    par = [-1] * N
    memo_val = [0] * N
    memo_pval = [0] * N
    cnt = 0
    st = [0] * (3 * N + 5)
    ans = 0
    while cnt >= 0:
        tmp = st[cnt]
        cnt -= 1
        (t, cur) = (tmp // M, tmp % M)
        a = A[cur]
        if t == 0:
            cnt += 1
            st[cnt] = 2 * M + cur
            ans += val[a] + pval[a]
            memo_val[cur] = val[a]
            memo_pval[cur] = pval[a]
            pval[a] = 1
            val[a] = 0
            i = he[cur]
            while i:
                nxt = to[i]
                i = ne[i]
                if nxt == par[cur]:
                    continue
                par[nxt] = cur
                cnt += 1
                st[cnt] = M + cur
                cnt += 1
                st[cnt] = nxt
        elif t == 1:
            pval[a] = 1
            val[a] = 0
        elif t == 2:
            val[a] = memo_val[cur] + 1
            pval[a] = memo_pval[cur]
    print(ans)