Introduction

This article presents various implementations of the DSU on Trees technique, also known as 'DSU on Tree'. The method involves processing each node's subtree in a specific order to efficiently compute answers for queries related to subtree properties.

Standard DSU Approach

The standard implementation uses a single array f to track color counts within the current subtree. When traversing a node, all light children's contributions are computed first, and their data is cleared from f. Then, the heavy child's contribution is processed and stored in f, preserving its state. Finally, all light children are processed again by iterating over each node in their subtrees and updating f accordingly.

To maintain efficiency, we track the maximum frequency (nowmx) and corresponding answer (nowans) during updates instead of scanning the entire f array. This ensures an overall complexity of O(n log n).

Complexity Analysis

Each node is visited at most log n times due to the number of light edges along its path to the root. Therefore, the time complexity is O(n log n), and space complexity remains O(n).

Implementation

#include <cstdio>
#include <iostream>
#define int long long
using namespace std;

const int MAXN = 100005;

int n;
int ans[MAXN], color[MAXN];
int heavy_child[MAXN], size[MAXN], freq[MAXN];
int visited[MAXN], current_value;
int max_count, result;
int head[MAXN], to[MAXN << 1], next_edge[MAXN << 1], edge_count = 2;

void add_edge(int x, int y) {
    next_edge[edge_count] = head[x];
    to[edge_count] = y;
    head[x] = edge_count++;
}

void initialize(int node, int parent) {
    size[node] = 1;
    for (int i = head[node]; i; i = next_edge[i]) {
        int child = to[i];
        if (child == parent) continue;
        initialize(child, node);
        size[node] += size[child];
        if (size[child] > size[heavy_child[node]]) {
            heavy_child[node] = child;
        }
    }
}

void dfs(int node, int parent, bool keep) {
    for (int i = head[node]; i; i = next_edge[i]) {
        int child = to[i];
        if (child == parent) continue;
        if (!keep && child == heavy_child[node]) continue;
        dfs(child, node, false);
    }

    if (keep) {
        if (visited[color[node]] < current_value) {
            freq[color[node]] = 0;
            visited[color[node]] = current_value;
        }
        if (++freq[color[node]] > max_count) {
            result = color[node];
            max_count = freq[color[node]];
        } else if (freq[color[node]] == max_count) {
            result += color[node];
        }
        return;
    }

    if (heavy_child[node]) {
        dfs(heavy_child[node], node, true);
    }

    for (int i = head[node]; i; i = next_edge[i]) {
        int child = to[i];
        if (child == parent || child == heavy_child[node]) continue;
        dfs(child, node, true);
    }

    if (visited[color[node]] < current_value) {
        freq[color[node]] = 0;
        visited[color[node]] = current_value;
    }
    if (++freq[color[node]] > max_count) {
        result = color[node];
        max_count = freq[color[node]];
    } else if (freq[color[node]] == max_count) {
        result += color[node];
    }

    ans[node] = result;
    if (!keep) {
        current_value++;
        result = 0;
        max_count = 0;
    }
}

signed main() {
    int x, y;
    scanf("%lld", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%lld", &color[i]);
    }
    for (int i = 1; i < n; i++) {
        scanf("%lld%lld", &x, &y);
        add_edge(x, y);
        add_edge(y, x);
    }
    initialize(1, 0);
    dfs(1, 0, false);
    for (int i = 1; i <= n; i++) {
        printf("%lld ", ans[i]);
    }
}

Using Hash Maps for Merging

An alternative approach employs hash maps to store color frequencies per node. Each node maintains a map of coler counts in its subtree. After processing children, the maps are merged into the parent's map.

For efficiency, we handle the heavy child separately, copying its map to the current node's map, then update it with the current node's color. Light children's maps are iterated through and merged into the current map.

Complexity Analysis

In worst-case scenarios where all colors are unique, each element is merged at most log n times, resulting in O(n log n) time complexity. Space complexity is O(n) as each node holds at most one copy of its subtere's data at any given time.

Implementation

#include <cstdio>
#include <unordered_map>
#define int long long
#define first fi
#define second se
using namespace std;

const int MAXN = 100005;

int head[MAXN], to[MAXN << 1], next_edge[MAXN << 1], edge_count = 2;
int n, color[MAXN];
int ans[MAXN], max_freq[MAXN];
int heavy_child[MAXN], size[MAXN];

void add_edge(int x, int y) {
    next_edge[edge_count] = head[x];
    to[edge_count] = y;
    head[x] = edge_count++;
}

unordered_map<int, int> color_map[MAXN];

void initialize(int node, int parent) {
    size[node] = 1;
    for (int i = head[node]; i; i = next_edge[i]) {
        int child = to[i];
        if (child == parent) continue;
        initialize(child, node);
        size[node] += size[child];
        if (size[child] > size[heavy_child[node]]) {
            heavy_child[node] = child;
        }
    }
}

void dfs(int node, int parent) {
    if (heavy_child[node]) {
        dfs(heavy_child[node], node);
        ans[node] = ans[heavy_child[node]];
        max_freq[node] = max_freq[heavy_child[node]];
        color_map[node] = color_map[heavy_child[node]];
        if ((++color_map[node][color[node]]) > max_freq[node]) {
            ans[node] = color[node];
            max_freq[node] = color_map[node][color[node]];
        } else if (color_map[node][color[node]] == max_freq[node]) {
            ans[node] += color[node];
        }
    } else {
        max_freq[node] = 1;
        ans[node] = color[node];
        color_map[node][color[node]] = 1;
        return;
    }

    for (int i = head[node]; i; i = next_edge[i]) {
        int child = to[i];
        if (child == parent || child == heavy_child[node]) continue;
        dfs(child, node);
        for (auto& entry : color_map[child]) {
            if ((color_map[node][entry.first] += entry.second) > max_freq[node]) {
                ans[node] = entry.first;
                max_freq[node] = color_map[node][entry.first];
            } else if (color_map[node][entry.first] == max_freq[node]) {
                ans[node] += entry.first;
            }
        }
        color_map[child].clear();
    }
}

signed main() {
    int x, y;
    scanf("%lld", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%lld", &color[i]);
    }
    for (int i = 1; i < n; i++) {
        scanf("%lld%lld", &x, &y);
        add_edge(x, y);
        add_edge(y, x);
    }
    initialize(1, 0);
    dfs(1, 0);
    for (int i = 1; i <= n; i++) {
        printf("%lld ", ans[i]);
    }
}

Heavy Chain Traversal

A third approach processes nodes along heavy chains, calculating answers from bottom to top. Each node is visited according to the number of heavy chains it belongs to, maintaining the same O(n log n) complexity. This method is less comonly used but conceptually equivalent to the previous approaches.