Path Queries (Codeforces 1213G)
Tags: Others, divide and conquer

You are given a weighted tree consisting of $$$n$$$ vertices. Recall that a tree is a connected graph without cycles. Vertices $$$u_i$$$ and $$$v_i$$$ are connected by an edge with weight $$$w_i$$$.You are given $$$m$$$ queries. The $$$i$$$-th query is given as an integer $$$q_i$$$. In this query you need to calculate the number of pairs of vertices $$$(u, v)$$$ ($$$u < v$$$) such that the maximum weight of an edge on a simple path between $$$u$$$ and $$$v$$$ doesn't exceed $$$q_i$$$.

Input: The first line of the input contains two integers $$$n$$$ and $$$m$$$ ($$$1 \le n, m \le 2 \cdot 10^5$$$) — the number of vertices in the tree and the number of queries.Each of the next $$$n - 1$$$ lines describes an edge of the tree. Edge $$$i$$$ is denoted by three integers $$$u_i$$$, $$$v_i$$$ and $$$w_i$$$ — the labels of vertices it connects ($$$1 \le u_i, v_i \le n$$$, $$$u_i \ne v_i$$$) and the weight of the edge ($$$1 \le w_i \le 2 \cdot 10^5$$$). It is guaranteed that the given edges form a tree.The last line of the input contains $$$m$$$ integers $$$q_1, q_2, \dots, q_m$$$ ($$$1 \le q_i \le 2 \cdot 10^5$$$), where $$$q_i$$$ is the maximum weight of an edge in the $$$i$$$-th query.

Output: Print $$$m$$$ integers — the answers to the queries. The $$$i$$$-th value should be equal to the number of pairs of vertices $$$(u, v)$$$ ($$$u < v$$$) such that the maximum weight of an edge on a simple path between $$$u$$$ and $$$v$$$ doesn't exceed $$$q_i$$$.Queries are numbered from $$$1$$$ to $$$m$$$ in the order of the input.

Note: The picture shows the tree from the first example: 

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

1 2
1 2

3 3
1 2 1
2 3 2
1 3 2

Sample Outputs:
21 7 15 21 3

0 0

1 3 3

Solution (Python):
def find_ancestor(i, father):
    if father[i] == i:
        return i
    father[i] = find_ancestor(father[i], father)
    return father[i]

def connect(i, j, father, n_child):
    i_anc = find_ancestor(i, father)
    j_anc = find_ancestor(j, father)
    if n_child[i_anc] > n_child[j_anc]:
        n_child[i_anc] += n_child[j_anc]
        father[j_anc] = i_anc
    else:
        n_child[j_anc] += n_child[i_anc]
        father[i_anc] = j_anc
(n, m) = map(int, input().split())
edges = []
father = [i for i in range(n)]
n_child = [1] * n
for i in range(n - 1):
    (i, j, w) = map(int, input().split())
    edges.append((i - 1, j - 1, w))
edges.sort(key=lambda x: -x[2])
queries = list(map(int, input().split()))
s_queries = sorted(queries)
ans = {}
w_limit = []
ans_cum = 0
for query in s_queries:
    while len(edges) and edges[-1][2] <= query:
        (i, j, w) = edges[-1]
        edges.pop()
        i_anc = find_ancestor(i, father)
        j_anc = find_ancestor(j, father)
        ans_cum += n_child[i_anc] * n_child[j_anc]
        connect(i, j, father, n_child)
    ans[query] = ans_cum
print(' '.join(list(map(str, [ans[query] for query in queries]))))

Solution (Python):
from collections import defaultdict
import sys
sys.setrecursionlimit(10 ** 6)
readline = sys.stdin.readline
(N, M) = map(int, readline().split())
E = []
for i in range(N - 1):
    (u, v, w) = map(int, readline().split())
    E.append((w, u - 1, v - 1))
E.sort()
Q = set()
MP = defaultdict(list)
for (i, q) in enumerate(map(int, readline().split())):
    MP[q].append(i)
    Q.add(q)
Q = list(Q)
Q.sort()

def fact(N):
    return N * fact(N - 1) % 100 if N > 1 else 1
fact(3000)

def root(x):
    if x == p[x]:
        return x
    y = x
    while y != p[y]:
        y = p[y]
    while x != y:
        (p[x], x) = (y, p[x])
    return y
(*p,) = range(N)
sz = [1] * N
c = 0

def unite(x, y):
    global c
    px = root(x)
    py = root(y)
    if px == py:
        return 0
    c += sz[px] * sz[py]
    if sz[px] < sz[py]:
        p[py] = px
        sz[px] += sz[py]
    else:
        p[px] = py
        sz[py] += sz[px]
    return 1
k = 0
ans = [N * (N - 1) // 2] * M
L = len(Q)
for (w, u, v) in E:
    while k < L and Q[k] < w:
        e = Q[k]
        for i in MP[e]:
            ans[i] = c
        k += 1
    unite(u, v)
sys.stdout.write(' '.join(map(str, ans)))
sys.stdout.write('\n')

Solution (Python):
from collections import defaultdict
from typing import DefaultDict
import heapq
import math
import sys
sys.setrecursionlimit(150000)

def nCr(n, r):
    return n * (n - 1) / 2

def findParent(i, parents):
    if parents[i] == i:
        return i
    ans = findParent(parents[i], parents)
    parents[i] = ans
    return ans

def merge(i, j, parents, count):
    f = findParent(i, parents)
    s = findParent(j, parents)
    if f == s:
        return f
    if count[f] >= count[s]:
        parents[s] = parents[f]
        count[f] += count[s]
        return f
    else:
        parents[f] = parents[s]
        count[s] += count[f]
        return s

def solve(n, tree, queries, weights):
    parents = [i for i in range(n + 1)]
    count = [1] * (n + 1)
    ans = {}
    res = 0
    weights.sort()
    for w in weights:
        for edge in tree[w]:
            f = findParent(edge[0], parents)
            s = findParent(edge[1], parents)
            (f, s) = (count[f], count[s])
            head = merge(edge[0], edge[1], parents, count)
            res += int(nCr(count[head], 2)) - int(nCr(f, 2)) - int(nCr(s, 2))
        ans[w] = res
    l = 0
    if len(weights) > 0:
        l = weights[-1]
    qs = [0] * l
    for i in range(1, l):
        if i in ans:
            qs[i] = ans[i]
        else:
            qs[i] = qs[i - 1]
    for q in queries:
        if q >= l:
            print(res, end=' ')
        else:
            print(qs[q], end=' ')

def main():
    (n, m) = map(int, input().split())
    tree = DefaultDict(list)
    weights = set()
    heap = []
    for _ in range(n - 1):
        (i, o, c) = map(int, input().split())
        tree[c].append([i, o])
        weights.add(c)
    queries = list(map(int, input().split()))
    weights = [i for i in weights]
    solve(n, tree, queries, weights)
main()