[算法学习] dsu on tree

简介

dsu on tree跟dsu没有关系，但是dsu on tree借鉴了dsu的启发式合并的思想。

它是用来解决一类树上的询问问题，一般这种问题有以下特征：

\(1.\)只有对子树的查询；

\(2.\)没有修改。

如果满足以上特征，那么dsu on tree很可能就可以派上用场了。

算法

我们以CF600E Lomsat gelral为例。

Descrption

一棵以\(1\)为根的树有 \(n\) 个结点，每个结点的颜色是\(c_i\)，每个颜色有一个编号，求树中每个子树的最多的颜色编号的和。

数据范围\(1\le n\le 10^5, 1\le c_i\le 10^5\)。

Solution

我们考虑暴力怎么写：遍历每一个节点 \(u\) ，然后把子树内的所有颜色暴力统计出来更新答案。

然后消除节点 \(u\) 的贡献，继续递归计算其他点的贡献。

复杂度\(O(n^2)\)，显然是很不优美的。

然后dsu on tree就登场了。它也是一个暴力，但是它结合了轻重链剖分，将复杂度降到了\(O(nlogn)\)。

我们先跑一个\(dfs1\)，记录节点\(u\)的重儿子\(heavy[u]\)。

接下来：

遍历每一个节点。
递归所有的轻儿子，递归进入时暴力加贡献，递归结束时暴力消除贡献。
递归重儿子，递归进入时暴力加贡献，递归结束时不消除贡献。
统计所有轻儿子的答案，并更新该节点的答案。
删除所有轻儿子对答案的影响。

复杂度\(O(nlogn)\)，可以通过本题。

Code

// Author: wlzhouzhuan

#pragma GCC optimize(2)

#pragma GCC optimize(3)

#include <bits/stdc++.h>

using namespace std;

#define ll long long

#define ull unsigned long long

#define rint register int

#define rep(i, l, r) for (rint i = l; i <= r; i++)

#define per(i, l, r) for (rint i = l; i >= r; i--)

#define mset(s, _) memset(s, _, sizeof(s))

#define pb push_back

#define pii pair <int, int>

#define mp(a, b) make_pair(a, b)

inline int read() {

  int x = 0, neg = 1; char op = getchar();

  while (!isdigit(op)) { if (op == '-') neg = -1; op = getchar(); }

  while (isdigit(op)) { x = 10 * x + op - '0'; op = getchar(); }

  return neg * x;

}

inline void print(int x) {

  if (x < 0) { putchar('-'); x = -x; }

  if (x >= 10) print(x / 10);

  putchar(x % 10 + '0');

}

const int N = 100005;

int n;

vector <int> adj[N];

void add(int u, int v) { adj[u].pb(v); }

int sz[N], fa[N], heavy[N]; // heavy[u] 表示u的重儿子

void dfs1(int u, int f) {

  sz[u] = 1, fa[u] = f;

  for (auto v: adj[u]) {

    if (v == f) continue;

    dfs1(v, u);

    sz[u] += sz[v];

    if (sz[v] > sz[heavy[u]]) {

      heavy[u] = v; // 更新重儿子

    }

  }

}

long long ans[N], sum;

int cnt[N], col[N], Son, Max;

void change(int u, int val) {

  cnt[col[u]] += val;

  if (cnt[col[u]] > Max) Max = cnt[col[u]], sum = col[u];

  else if (cnt[col[u]] == Max) sum += col[u];

  for (auto v: adj[u]) {

    // 由于重儿子的信息没有被删去，所以已经统计过了，不能再计算

    if (v == fa[u] || v == Son) continue;

    change(v, val);

  }

}

void dfs2(int u, int keep) {

  for (auto v: adj[u]) {

    if (v == fa[u] || v == heavy[u]) continue;

    dfs2(v, 0); // 遍历u的轻儿子

  }

  if (heavy[u]) dfs2(heavy[u], 1), Son = heavy[u];

  change(u, 1), ans[u] = sum, Son = 0;

  if (!keep) change(u, -1), sum = Max = 0;

}

int main() {

  n = read();

  for (rint i = 1; i <= n; i++) {

    col[i] = read();

  }

  for (rint i = 1; i < n; i++) {

    int x = read(), y = read();

    add(x, y), add(y, x);

  }

  int root = 1; // 题目默认1为根

  dfs1(root, 0);

  dfs2(root, 0);

  for (rint i = 1; i <= n; i++) {

    printf("%lld ", ans[i]);

  }

  return 0;

}

巴特西

[算法学习] dsu on tree

简介

算法

Descrption

Solution

Code

最新文章

热门文章