解決樹上詢問問題,沒有修改
時間復雜度:\(O(nlogn)\)
例題:https://codeforc.es/contest/600/problem/E
題意:給定一顆樹,每個節點有一個顏色,求出子樹中顏色最多的顏色值之和。
代碼:
#include<bits/stdc++.h>
using namespace std;
using LL = long long;
vector<int> color;
struct HLD{
vector<vector<int>> e;
vector<int> siz, son, cnt;
vector<LL> ans;
LL sum, Max;
int hson;
HLD(int n){
e.resize(n + 1);
siz.resize(n + 1);
son.resize(n + 1);
ans.resize(n + 1);
cnt.resize(n + 1);
hson = 0;
sum = 0;
Max = 0;
}
void add(int u, int v){
e[u].push_back(v);
e[v].push_back(u);
}
void dfs1(int u, int fa){
siz[u] = 1;
for (auto v : e[u]){
if (v == fa) continue;
dfs1(v, u);
siz[u] += siz[v];
if (siz[v] > siz[son[u]]) son[u] = v;
}
}
void calc(int u, int fa, int val){
cnt[color[u]] += val;
if (cnt[color[u]] > Max){
Max = cnt[color[u]];
sum = color[u];
}
else if (cnt[color[u]] == Max){
sum += color[u];
}
for (auto v : e[u]){
if (v == fa || v == hson) continue;
calc(v, u, val);
}
}
void dfs2(int u, int fa, int op){
for (auto v : e[u]){
if (v == fa || v == son[u]) continue;
dfs2(v, u, 0);
}
if (son[u]){
dfs2(son[u], u, 1);
hson = son[u]; //記錄重鏈編號,計算的時候跳過
}
calc(u, fa, 1);
hson = 0; //消除的時候所有兒子都清除
ans[u] = sum;
if (!op){
calc(u, fa, -1);
sum = 0;
Max = 0;
}
}
};
int main(){
ios::sync_with_stdio(false);cin.tie(0);
int n;
cin >> n;
color.resize(n + 1);
for (int i = 1; i <= n; i ++ ){
cin >> color[i];
}
HLD t(n);
for (int i = 0; i < n - 1; i ++ ){
int u, v;
cin >> u >> v;
t.add(u, v);
}
t.dfs1(1, 0);
t.dfs2(1, 0, 0);
for (int i = 1; i <= n; i ++ ){
cout << t.ans[i] << " \n"[i == n];
}
return 0;
}
浙公網安備 33010602011771號