nowcoder 79F 小H和圣诞树 换根 DP + 根号分治

设节点个数大于 $\sqrt n$ 的颜色为关键颜色，那么可以证明关键颜色最多有 $\sqrt n$ 个.
对于每个关键颜色，暴力预处理出该颜色到查询中另一个颜色的距离和.
对于不是关键颜色的询问，直接建立虚树进行统计即可.
由于不是关键颜色，节点数最多为 $\sqrt n$ ，那么时间复杂度是 $O(2\times n\sqrt n)$.
总时间复杂度为 $O(n\sqrt n)$，这个就叫做根号分治.

#include <cstdio>
#include <algorithm>
#include <vector>
#include <cmath>
#include <map>
#define N 100003
#define ll long long
#define setIO(s) freopen(s".in", "r" , stdin)  , freopen(s".out", "w" , stdout)
using namespace std;
namespace IO
{
    char *p1,*p2,buf[100000];
    #define nc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++)
    int readint() {int x=0; char c=nc(); while(c<48) c=nc(); while(c>47) x=(((x<<2)+x)<<1)+(c^48),c=nc(); return x;}
    ll readll() {ll x=0; char c=nc(); while(c<48) c=nc(); while(c>47) x=(((x<<2)+x)<<1)+(c^48),c=nc(); return x;}
};
vector <int> G[N], ty[N], node;
int n , edges, tim, toop;
ll dis[N], depth[N];
int col[N], tax[N], id[N], A[N], size[N], S[N];
int hd[N], nex[N << 1], to[N << 1], top[N], dfn[N], fa[N], dep[N], son[N], siz[N];
ll val[N << 1];
bool cmp(int a, int b)
{
    return dfn[a] < dfn[b];
}
inline void addedge(int u, int v, int c)
{
    nex[++ edges] = hd[u], hd[u] = edges, to[edges] = v, val[edges] = 1ll * c;
}
void dfs1(int u, int ff)
{
    int i, v;
    fa[u] = ff, dep[u] = dep[ff] + 1, dfn[u] = ++ tim, siz[u] = 1;
    for(i = hd[u] ; i ; i = nex[i])
    {
        v = to[i];
        if(v == ff) continue;
        depth[v] = depth[u] + 1ll * val[i], dfs1(v, u), siz[u] += siz[v];
        if(siz[v] > siz[son[u]]) son[u] = v;
    }
}
void dfs2(int u, int tp)
{
    top[u] = tp;
    if(son[u]) dfs2(son[u], tp);
    for(int i = hd[u] ; i ; i = nex[i])
    {
        int v = to[i];
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}
inline int LCA(int x, int y)
{
    while(top[x] ^ top[y])
    {
        dep[top[x]] > dep[top[y]] ? x = fa[top[x]] : y = fa[top[y]];
    }
    return dep[x] < dep[y] ? x : y;
}
inline ll Dis(int x, int y)
{
    return depth[x] + depth[y] - (depth[LCA(x, y)] << 1);
}
void solve1(int u, int ff, int cur)
{
    size[u] = (col[u] == cur), dis[u] = 0;
    for(int i = hd[u] ; i ; i = nex[i])
    {
        int v = to[i];
        if(v == ff) continue;
        solve1(v, u, cur),size[u] += size[v], dis[u] += (dis[v] + 1ll * size[v] * val[i]);
    }
}
void solve(int u, int ff, int cur)
{
    for(int i = hd[u] ; i ; i = nex[i])
    {
        int v = to[i];
        if(v == ff) continue;
        dis[v] += (dis[u] - dis[v] - 1ll * size[v] * val[i] + 1ll * (tax[cur] - size[v]) * val[i]);
        solve(v, u, cur);
    }
}
inline void addvir(int u, int v)
{
    G[u].push_back(v);
}
inline void insert(int x)
{
    if(toop < 2)
    {
        S[++ toop] = x;
        return ;
    }
    int lca = LCA(x, S[toop]);
    if(lca != S[toop])
    {
        while(toop > 1 && dep[S[toop - 1]] >= dep[lca]) addvir(S[toop - 1], S[toop]),-- toop;
        if(S[toop] != lca) addvir(lca, S[toop]), S[toop] = lca;
    }
    S[++ toop] = x;
}
void pre(int u, int ff, int cur)
{
    size[u] = (col[u] == cur), dis[u] = 0;
    for(int i = 0; i < G[u].size(); ++ i)
    {
        int v = G[u][i];
        pre(v, u, cur), size[u] += size[v], dis[u] += dis[v] + 1ll * size[v] * Dis(v, u);
    }
}
void work(int u, int ff, int cur)
{
    for(int i = 0; i < G[u].size() ; ++ i)
    {
        int v = G[u][i];
        dis[v] += (dis[u] - dis[v] - 1ll * size[v] * Dis(u, v) + 1ll * (tax[cur] - size[v]) * Dis(u, v));
        work(v, u, cur);
    }
}
void clear(int u)
{
    size[u] = dis[u] = 0;
    for(int i = 0; i < G[u].size(); ++ i) clear(G[u][i]) ;
    G[u].clear();
}
struct Node
{
    int a, b;
}ask[N];
vector < int > P[N];
vector < ll > answer[N];
int point[N];
int main()
{
    using namespace IO;
    // setIO("input");
    int i , j, idx = 0, m, Q;
    n = readint();
    m = sqrt(n);
    for(i = 1; i <= n ; ++ i) col[i] = readint(), ++tax[col[i]], ty[col[i]].push_back(i);
    for(i = 1; i < n ; ++ i)
    {
        int a = readint(), b = readint(), c = readint();
        addedge(a, b, c), addedge(b, a, c);
    }
    dfs1(1, 0), dfs2(1, 1);
    for(i = 1; i <= n ; ++ i)   if(tax[i] >= m) id[i] = ++idx;
    Q = readint();
    for(i = 1; i <= Q; ++ i)
    {
        ask[i].a = readint(), ask[i].b = readint();
        if(tax[ask[i].a] < tax[ask[i].b]) swap(ask[i].a, ask[i].b);
        if(tax[ask[i].a] >= m) P[ask[i].a].push_back(ask[i].b);
    }
    for(i = 1; i <= n ; ++ i)
    {
        if(tax[i] >= m)
        {
            solve1(1, 0, i), solve(1, 0, i);
            for(j = 0 ; j < P[i].size() ; ++ j)
            {
                int cur = P[i][j];
                ll re = 0;
                for(int k = 0; k < ty[cur].size(); ++ k)
                {
                    re += dis[ty[cur][k]];
                }
                answer[i].push_back(re);
            }
        }
    }
    for(int cas = 1; cas <= Q; ++ cas)
    {
        int a, b;
        a = ask[cas].a, b = ask[cas].b;
        if(tax[a] >= m) printf("%lld\n", a == b ? answer[a][point[a] ++ ] / 2 : answer[a][point[a] ++ ]);
        else
        {
            int tmp = 0;
            ll re = 0;
            for(i = 0; i < ty[a].size(); ++ i) A[++ tmp] = ty[a][i];
            for(i = 0; i < ty[b].size(); ++ i) A[++ tmp] = ty[b][i];
            sort(A + 1, A + 1 + tmp, cmp);
            tmp = unique(A + 1, A + 1 + tmp) - (A + 1);
            toop = 0;
            if(A[1] != 1) S[++ toop] = 1;
            for(i = 1 ; i <= tmp ; ++ i) insert(A[i]);
            while(toop > 1) addvir(S[toop - 1], S[toop]), --toop;
            pre(1, 0, b), work(1, 0, b);
            for(i = 0; i < ty[a].size(); ++ i) re += dis[ty[a][i]];
            printf("%lld\n", a == b ? re / 2 : re);
        }
    }
    return 0;
}