DSU on Tree
(tree/dsu_on_tree.cpp)

• View this file on GitHub
• Last update: 2026-03-13 21:29:59+09:00

説明

各頂点を根とする部分木について、追加・削除可能なデータ構造を使ってクエリを処理する。 全体計算量は、update と clear が $O(1)$ なら $O(N log N)$。

できること

• DSUonTree<G> dsu(g, root) 根 root を持つ木 g を前処理する
• int idx(int v) Euler Tour での v の位置を返す
• int begin(int v) 部分木 v の Euler Tour 区間の左端を返す
• int end(int v) 部分木 v の Euler Tour 区間の右端を返す
• void run(update, query, clear, reset) update(v) で頂点を追加し、query(v) で部分木 v の答えを確定する keep == false の片付け時には、部分木内の各頂点に clear(v) を呼んだあと reset() を呼ぶ

使い方

update は頂点 1 個を現在のデータ構造へ反映する。 query は「現在のデータ構造がちょうど頂点 v の部分木を表している」タイミングで呼ばれる。 削除が重い場合は clear を空にして、reset で全体を初期化してもよい。

vector<int> ans(n);
auto update = [&](int v) {
    // 頂点 v を追加
};
auto query = [&](int v) {
    // 部分木 v の答えを記録
};
auto clear = [&](int v) {
    // 頂点 v を削除
};
auto reset = [&]() {
    // 全体を初期状態へ戻す
};

DSUonTree dsu(g, 0);
dsu.run(update, query, clear, reset);

実装上の補足

g[v][0] が heavy child になるように隣接順を並べ替えるので、g は破壊的に変更される。

Verified with

• test/yosupo_vertex_add_subtree_sum_dsu_on_tree.test.cpp

Code

template<class G>
struct DSUonTree {
    G &g;
    int n, root, ord;
    vector<int> sub_size, euler, down, up;

    explicit DSUonTree(G &g, int root = 0)
        : g(g), n(g.size()), root(root), ord(0),
          sub_size(n), euler(n), down(n), up(n) {
        dfs_sz(root, -1);
        dfs_euler(root, -1);
    }

    int idx(int v) const {
        return down[v];
    }

    int begin(int v) const {
        return down[v];
    }

    int end(int v) const {
        return up[v];
    }

    template<class UPDATE, class QUERY, class CLEAR, class RESET>
    void run(UPDATE update, QUERY query, CLEAR clear, RESET reset) {
        auto dfs = [&](auto &&self, int v, int p, bool keep) -> void {
            int heavy = (g[v].empty() || g[v][0] == p ? -1 : g[v][0]);
            for (auto &&u : g[v]) {
                if (u == p || u == heavy) continue;
                self(self, u, v, false);
            }
            if (heavy != -1) self(self, heavy, v, true);
            for (auto &&u : g[v]) {
                if (u == p || u == heavy) continue;
                for (int i = down[u]; i < up[u]; ++i) update(euler[i]);
            }
            update(v);
            query(v);
            if (!keep) {
                for (int i = down[v]; i < up[v]; ++i) clear(euler[i]);
                reset();
            }
        };
        dfs(dfs, root, -1, false);
    }

private:
    void dfs_sz(int v, int p) {
        sub_size[v] = 1;
        int heavy_idx = -1;
        for (int i = 0; i < (int)g[v].size(); ++i) {
            int u = g[v][i];
            if (u == p) continue;
            dfs_sz(u, v);
            sub_size[v] += sub_size[u];
            if (heavy_idx == -1 || sub_size[u] > sub_size[g[v][heavy_idx]]) {
                heavy_idx = i;
            }
        }
        if (heavy_idx > 0) swap(g[v][0], g[v][heavy_idx]);
    }

    void dfs_euler(int v, int p) {
        down[v] = ord;
        euler[ord++] = v;
        for (auto &&u : g[v]) {
            if (u == p) continue;
            dfs_euler(u, v);
        }
        up[v] = ord;
    }
};

/**
 * @brief DSU on Tree
 */

#line 1 "tree/dsu_on_tree.cpp"
template<class G>
struct DSUonTree {
    G &g;
    int n, root, ord;
    vector<int> sub_size, euler, down, up;

    explicit DSUonTree(G &g, int root = 0)
        : g(g), n(g.size()), root(root), ord(0),
          sub_size(n), euler(n), down(n), up(n) {
        dfs_sz(root, -1);
        dfs_euler(root, -1);
    }

    int idx(int v) const {
        return down[v];
    }

    int begin(int v) const {
        return down[v];
    }

    int end(int v) const {
        return up[v];
    }

    template<class UPDATE, class QUERY, class CLEAR, class RESET>
    void run(UPDATE update, QUERY query, CLEAR clear, RESET reset) {
        auto dfs = [&](auto &&self, int v, int p, bool keep) -> void {
            int heavy = (g[v].empty() || g[v][0] == p ? -1 : g[v][0]);
            for (auto &&u : g[v]) {
                if (u == p || u == heavy) continue;
                self(self, u, v, false);
            }
            if (heavy != -1) self(self, heavy, v, true);
            for (auto &&u : g[v]) {
                if (u == p || u == heavy) continue;
                for (int i = down[u]; i < up[u]; ++i) update(euler[i]);
            }
            update(v);
            query(v);
            if (!keep) {
                for (int i = down[v]; i < up[v]; ++i) clear(euler[i]);
                reset();
            }
        };
        dfs(dfs, root, -1, false);
    }

private:
    void dfs_sz(int v, int p) {
        sub_size[v] = 1;
        int heavy_idx = -1;
        for (int i = 0; i < (int)g[v].size(); ++i) {
            int u = g[v][i];
            if (u == p) continue;
            dfs_sz(u, v);
            sub_size[v] += sub_size[u];
            if (heavy_idx == -1 || sub_size[u] > sub_size[g[v][heavy_idx]]) {
                heavy_idx = i;
            }
        }
        if (heavy_idx > 0) swap(g[v][0], g[v][heavy_idx]);
    }

    void dfs_euler(int v, int p) {
        down[v] = ord;
        euler[ord++] = v;
        for (auto &&u : g[v]) {
            if (u == p) continue;
            dfs_euler(u, v);
        }
        up[v] = ord;
    }
};

/**
 * @brief DSU on Tree
 */

Back to top page