Codeforces 600E - Lomsat gelral 「$Dsu on tree$模板」

Posted 2021-02-08 colythme

With $Dsu on tree$ we can answer queries of this type:

How many vertices in the subtree of vertex $v$ has some property in $O (n log n)$ time (for all of the queries)?

这题写的是轻重儿子（重链剖分）版本的 $Dsu on tree$

具体流程如下：

每次先递归计算轻儿子，再单独递归重儿子，计算完后轻儿子的一些信息需要删掉，但是重儿子的信息无需删除，如此出解，相当于是优化了暴力的多余部分

每个节点会作为轻儿子被计算，重链剖分上垂直有 $log n$ 条链，故复杂度 $O (n log n)$

代码

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <cstring>
  4 
  5 using namespace std;
  6 
  7 typedef long long LL;
  8 
  9 const int MAXN = 1e05 + 10;
 10 const int MAXM = 1e05 + 10;
 11 const int MAXC = 1e05 + 10;
 12 
 13 struct LinkedForwardStar {
 14     int to;
 15 
 16     int next;
 17 } ;
 18 
 19 LinkedForwardStar Link[MAXM << 1];
 20 int Head[MAXN]= {0};
 21 int size = 0;
 22 
 23 void Insert (int u, int v) {
 24     Link[++ size].to = v;
 25     Link[size].next = Head[u];
 26 
 27     Head[u] = size;
 28 }
 29 
 30 int N;
 31 int colour[MAXN];
 32 
 33 int son[MAXN]= {0};
 34 int subsize[MAXN]= {0};
 35 void DFS (int root, int father) {
 36     son[root] = - 1;
 37     subsize[root] = 1;
 38     for (int i = Head[root]; i; i = Link[i].next) {
 39         int v = Link[i].to;
 40         if (v == father)
 41             continue;
 42         DFS (v, root);
 43         subsize[root] += subsize[v];
 44         if (son[root] == - 1 || subsize[v] > subsize[son[root]])
 45             son[root] = v;
 46     }
 47 }
 48 int vis[MAXN]= {0};
 49 int total[MAXC]= {0};
 50 int maxv = 0;
 51 LL sum = 0;
 52 void calc (int root, int father, int delta) { // 统计答案
 53     total[colour[root]] += delta;
 54     if (delta > 0 && total[colour[root]] >= maxv) {
 55         if (total[colour[root]] > maxv)
 56             sum = 0, maxv = total[colour[root]];
 57         sum += colour[root];
 58     }
 59     for (int i = Head[root]; i; i = Link[i].next) {
 60         int v = Link[i].to;
 61         if (v == father || vis[v])
 62             continue;
 63         calc (v, root, delta);
 64     }
 65 }
 66 LL answer[MAXN]= {0};
 67 void Solve (int root, int father, int type) { // type表示是不是重儿子信息
 68     for (int i = Head[root]; i; i = Link[i].next) {
 69         int v = Link[i].to;
 70         if (v == father || v == son[root])
 71             continue;
 72         Solve (v, root, 0);
 73     }
 74     if (~ son[root])
 75         Solve (son[root], root, 1), vis[son[root]] = 1;
 76     calc (root, father, 1);
 77     answer[root] = sum;
 78     if (~ son[root])
 79         vis[son[root]] = 0;
 80     if (! type) // 如果是轻儿子信息就需删除
 81         calc (root, father, - 1), maxv = sum = 0;
 82 }
 83 
 84 int getnum () {
 85     int num = 0;
 86     char ch = getchar ();
 87 
 88     while (! isdigit (ch))
 89         ch = getchar ();
 90     while (isdigit (ch))
 91         num = (num << 3) + (num << 1) + ch - ‘0‘, ch = getchar ();
 92 
 93     return num;
 94 }
 95 
 96 int main () {
 97     N = getnum ();
 98     for (int i = 1; i <= N; i ++)
 99         colour[i] = getnum ();
100     for (int i = 1; i < N; i ++) {
101         int u = getnum (), v = getnum ();
102         Insert (u, v), Insert (v, u);
103     }
104     DFS (1, 0), Solve (1, 0, 0);
105     for (int i = 1; i <= N; i ++) {
106         if (i > 1)
107             putchar (‘ ‘);
108         printf ("%lld", answer[i]);
109     }
110     puts ("");
111 
112     return 0;
113 }
114 
115 /*
116 4
117 1 2 3 4
118 1 2
119 2 3
120 2 4
121 */
122 
123 /*
124 15
125 1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
126 1 2
127 1 3
128 1 4
129 1 14
130 1 15
131 2 5
132 2 6
133 2 7
134 3 8
135 3 9
136 3 10
137 4 11
138 4 12
139 4 13
140 */