树的果实(赛氪冬季赛)【启发式合并 dsu on tree】

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

传送门

题意

给定一颗有根树（1为根），$m$次询问，当前询问的子树中的价值是多少？ 价值定义为，取子树中所有边权，$\sum_{w}{(w*cnt_w)^2},w\in子树边权集合$，$cnt_w$ 为 $w$出现的次数

分析

题意很清楚，求的是子树信息统计，此时我们可以使用树上启发式合并来解决（优美的暴力）在之前有一篇文章有写【dsu个人理解】

暴力的写法是，对于每个子树都进行统计，正确性保证，时间复杂度为$On^2$，但是通过启发式合并，能够降到$Onlogn$

细节： 里面有统计边权出现的次数的需求，对于轻儿子需要删除统计信息，此时如果每次都重新清空记录的话，时间复杂度会高，可以使用记录版本信息进行优化，见结构体 $Node$ 的实现

代码

//D
/*
  @Author: YooQ
*/
#include <bits/stdc++.h>
using namespace std;
#define sc scanf
#define pr printf
#define ll long long
#define int long long
#define FILE_OUT freopen("out", "w", stdout);
#define FILE_IN freopen("in", "r", stdin);
#define debug(x) cout << #x << ": " << x << "\n";
#define CIN(arr) for (int i = 1; i <= N; ++i) cin >> arr[i];
#define max3(a, b, c) max(a, max(b, c))
#define min3(a, b, c) min(a, min(b, c))
#define MAX(a, b) (a >= b ? a : a = b)
#define MIN(a, b) (a <= b ? a : a = b)
#define AC 0
#define WA 1
#define INF 0x3f3f3f3f
const ll MAX_N = 1e6+5;
const ll MOD = 1e9+7;
int N, M, K;

int arr[MAX_N];

int head[MAX_N];
int tot = 0;
struct Edge {
	int to, nxt, w;
}edge[MAX_N];

void addEdge(int u, int v, int w) {
	edge[tot].nxt = head[u];
	edge[tot].to = v;
	edge[tot].w = w;
	head[u] = tot++;
	edge[tot].nxt = head[v];
	edge[tot].to = u;
	edge[tot].w = w;
	head[v] = tot++;
}

int sz[MAX_N];
int son[MAX_N];
int _son[MAX_N];

void dfs1(int u, int from) {
	sz[u] = 1;
	son[u] = 0;
	_son[u] = 0;
	int v;
	for (int i = head[u]; ~i; i = edge[i].nxt) {
		if ((v = edge[i].to) == from) continue;
		dfs1(v, u);
		sz[u] += sz[v];
		if (sz[v] > sz[son[u]]) {
			son[u] = v;
			_son[u] = i;
		}
	}
}

int ans[MAX_N];

struct Node {
	int vis[MAX_N];
	int version[MAX_N];
	int ver = 1;
	int res = 0;
	
	void add(int x) {
		if (version[x] != ver) {
			version[x] = ver;
			vis[x] = 1;
			res += x*x;
		} else {
			res -= vis[x]*x*vis[x]*x;
			++vis[x];
			res += vis[x]*x*vis[x]*x;
		}
	}
	
	void clear() {
		++ver;
		res = 0;
	}
	
}cnt;

void calc(int u, int from, int cut) {
	int v;
	for (int i = head[u]; ~i; i = edge[i].nxt) {
		if ((v = edge[i].to) == from || v == cut) continue;
		cnt.add(edge[i].w);
		calc(v, u, cut);
	}
}

void dfs(int u, int from, int del) {
	int v;
	for (int i = head[u]; ~i; i = edge[i].nxt) {
		if ((v = edge[i].to) == from || v == son[u]) continue;
		dfs(v, u, 1);
	}
	if (son[u]) dfs(son[u], u, 0), cnt.add(edge[_son[u]].w);
	calc(u, from, son[u]);
	ans[u] = cnt.res;
	if (del) cnt.clear();
}

void init() {
	memset(head, -1, sizeof head);
	tot = 0;
}

void solve() {
	init();
    cin >> N >> K >> M;
    int u, v, w;
    for (int i = 1; i <= N-1; ++i) {
    	cin >> u >> v >> w;
    	addEdge(u, v, w);
	}
    dfs1(1, 0);
    dfs(1, 0, 0);
    
    for (int i = 1; i <= K; ++i) {
    	cin >> u;
    	cout << ans[u] << "\n";
	}
}

signed main() {
	#ifndef ONLINE_JUDGE
//	FILE_IN
	FILE_OUT
	#endif
	int T = 1;//cin >> T;
	while (T--) solve();

	return AC;
}