旅行

题目大意

给定一棵$n(1 \le n \le 2 \times 10^5)$个节点的树，每个结点都有一个类型$c_i(1 \le c_i \le n)$和权值$w_i(1 \le w_i \le n)$，定义一条链有权值当且仅当起始点的类型相同，且权值为起始点的权值和。求取任意条不相交的链的权值和最大的值。

解题思路

对于每个节点，考虑有一条链覆盖它和没有链覆盖它对应的最大的子树权值和。没有链覆盖时，直接将所有子树的$dp$值相加即可；被链覆盖时，考虑用线段树合并维护每种类型所对应的链上的子树的$dp$值的和。在进行线段树的合并的时候，是对两棵树的节点取最值，这样就成功维护了$dp$值。

参考代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#include <bits/stdc++.h>
#define maxn 200100
#define int long long
using namespace std;
const double eps = 1e-8;
struct node{
    int s0, s1, cnt, add;
}tr[20 * maxn];
int c[maxn], w[maxn], f[maxn], root[maxn], idx, n, tt;
vector<int> e[maxn];

void pushdown(int u) {
    int tmp = tr[u].add;
    if(tr[u].s0)  tr[tr[u].s0].cnt += tmp, tr[tr[u].s0].add += tmp;
    if(tr[u].s1)  tr[tr[u].s1].cnt += tmp, tr[tr[u].s1].add += tmp;
    tr[u].add = 0;
}

void merge(int &u, int v, int l, int r, int& res) {
    if(!u || !v) {
        u = u + v;
        return ;
    }
    pushdown(u);  pushdown(v);
    if(l == r) {
        res = max(res, tr[u].cnt + tr[v].cnt + tt);
        tr[u].cnt = max(tr[u].cnt, tr[v].cnt);
    }
    else {
        int mid = l + r >> 1;
        merge(tr[u].s0, tr[v].s0, l, mid, res);
        merge(tr[u].s1, tr[v].s1, mid + 1, r, res);
    }
}

void modify(int &u, int l, int r, int x, int d) {
    if(!u)  u = ++idx;
    if(l == r) {
        tr[u].cnt = d;
        return ;
    }
    int mid = l + r >> 1;
    if(x <= mid)  modify(tr[u].s0, l, mid, x, d);
    if(x > mid)  modify(tr[u].s1, mid + 1, r, x, d);
}

void dfs(int x, int fa) {
    modify(root[x], 1, n, c[x], w[x]);
    int sum = 0;
    for (auto u : e[x]) {
        if(u == fa)  continue;
        dfs(u, x);
        sum += f[u];
    }
    f[x] = sum;
    tt = sum;
    for (auto u : e[x]) {
        if(u == fa)  continue;
        tr[root[u]].cnt -= f[u];
        tr[root[u]].add -= f[u];
        merge(root[x], root[u], 1, n, f[x]);
    }
    tr[root[x]].cnt += sum;
    tr[root[x]].add += sum;
}

void solve() {
    cin >> n;
    for (int i = 1; i <= n; ++i)  cin >> c[i];
    for (int i = 1; i <= n; ++i)  cin >> w[i];
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        e[u].push_back(v);  e[v].push_back(u);
    }
    dfs(1, -1);
    // for (int i = 1; i <= n; ++i)  cout << f[i] << '\n';
    cout << f[1] << '\n';
    for (int i = 1; i <= idx; ++i)  tr[i].s0 = tr[i].s1 = tr[i].add = tr[i].cnt = 0;
    for (int i = 1; i <= n; ++i)  f[i] = root[i] = 0, e[i].clear(), idx = 0;
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);  cout.tie(0);
    int t = 1;  cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}