sotanishy's competitive programming library

sotanishy's code snippets for competitive programming

View the Project on GitHub sotanishy/cp-library-cpp

:heavy_check_mark: Range Contour Aggregation
(tree/range_contour_aggregation.hpp)

Description

木上の等高線集約クエリ (の変種1) (参考文献を参照) を扱う．

空間計算量: $O(n\log n)$

Operations

Reference

Depends on

Verified with

Code

#pragma once
#include <map>
#include <queue>
#include <utility>
#include <vector>

#include "../data-structure/fenwick_tree.hpp"
#include "centroid_decomposition.hpp"

template <typename Group>
class RangeContourAggregation {
    using T = Group::T;

   public:
    RangeContourAggregation() = default;
    explicit RangeContourAggregation(const std::vector<std::vector<int>>& G)
        : seg(G.size()),
          seg_sub(G.size()),
          dist(G.size()),
          idx(G.size()),
          sub(G.size()),
          idx_sub(G.size()),
          first(G.size()),
          first_sub(G.size()) {
        std::tie(level, sz_comp, par) = centroid_decomposition(G);

        for (int v = 0; v < (int)G.size(); ++v) {
            // calculate for each vertex u,
            // dist: dist from v
            // idx: index in the BFS order
            // sub: subtree of v that u belongs to
            // idx_sub: index in the subtree

            // calculate for each depth d,
            // first: first index of depth d
            // first_sub: first index of depth d in the subtree i

            dist[v][v] = 0;
            std::queue<std::pair<int, int>> que;
            que.push({v, -1});
            int k = 0, sub_idx = 0;
            std::vector<int> ks;

            while (!que.empty()) {
                auto [u, i] = que.front();
                que.pop();

                // update component info
                int d = dist[v][u];
                if (d >= (int)first[v].size()) {
                    first[v].push_back(k);
                }
                idx[v][u] = k++;
                // enter subtree
                if (d == 1) {
                    i = sub_idx++;
                    ks.push_back(0);
                    first_sub[v].emplace_back();
                }
                // update subtree info
                if (i != -1) {
                    sub[v][u] = i;
                    if (d - 1 >= (int)first_sub[v][i].size()) {
                        first_sub[v][i].push_back(ks[i]);
                    }
                    idx_sub[v][u] = ks[i]++;
                }

                for (int w : G[u]) {
                    if (level[w] > level[v] && !dist[v].contains(w)) {
                        dist[v][w] = d + 1;
                        que.push({w, i});
                    }
                }
            }

            first[v].push_back(k);
            seg[v] = FenwickTree<Group>(k);
            for (int i = 0; i < sub_idx; ++i) {
                first_sub[v][i].push_back(ks[i]);
                seg_sub[v].emplace_back(ks[i]);
            }
        }
    }

    void update(int v, T x) {
        seg[v].update(0, x);
        for (int p = par[v]; p != -1; p = par[p]) {
            seg[p].update(idx[p][v], x);
            int i = sub[p][v];
            seg_sub[p][i].update(idx_sub[p][v], x);
        }
    }

    T fold(int v, int d) const {
        int dd = std::min((int)first[v].size() - 1, d);
        T res = seg[v].prefix_fold(first[v][dd]);
        for (int p = par[v]; p != -1; p = par[p]) {
            int dep = d - dist[p].at(v);
            if (dep >= 0) {
                dd = std::min((int)first[p].size() - 1, dep);
                res = Group::op(res, seg[p].prefix_fold(first[p][dd]));
                if (dd > 0) {
                    int i = sub[p].at(v);
                    dd = std::min((int)first_sub[p][i].size() - 1, dep - 1);
                    res = Group::op(res, Group::inv(seg_sub[p][i].prefix_fold(
                                             first_sub[p][i][dd])));
                }
            }
        }
        return res;
    }

   private:
    std::vector<int> level, sz_comp, par;
    std::vector<FenwickTree<Group>> seg;
    std::vector<std::vector<FenwickTree<Group>>> seg_sub;
    std::vector<std::unordered_map<int, int>> dist, idx, sub, idx_sub;
    std::vector<std::vector<int>> first;
    std::vector<std::vector<std::vector<int>>> first_sub;
};
#line 2 "tree/range_contour_aggregation.hpp"
#include <map>
#include <queue>
#include <utility>
#include <vector>

#line 2 "data-structure/fenwick_tree.hpp"
#include <functional>
#line 4 "data-structure/fenwick_tree.hpp"

template <typename M>
class FenwickTree {
    using T = M::T;

   public:
    FenwickTree() = default;
    explicit FenwickTree(int n) : n(n), data(n + 1, M::id()) {}

    T prefix_fold(int i) const {
        T ret = M::id();
        for (; i > 0; i -= i & -i) ret = M::op(ret, data[i]);
        return ret;
    }

    void update(int i, const T& x) {
        for (++i; i <= n; i += i & -i) data[i] = M::op(data[i], x);
    }

    int lower_bound(const T& x) const { return lower_bound(x, std::less<>()); }

    template <typename Compare>
    int lower_bound(const T& x, Compare cmp) const {
        if (!cmp(M::id(), x)) return 0;
        int k = 1;
        while (k * 2 <= n) k <<= 1;
        int i = 0;
        T v = M::id();
        for (; k > 0; k >>= 1) {
            if (i + k > n) continue;
            T nv = M::op(v, data[i + k]);
            if (cmp(nv, x)) {
                v = nv;
                i += k;
            }
        }
        return i + 1;
    }

   private:
    int n;
    std::vector<T> data;
};
#line 2 "tree/centroid_decomposition.hpp"
#include <tuple>
#line 4 "tree/centroid_decomposition.hpp"

std::tuple<std::vector<int>, std::vector<int>, std::vector<int>>
centroid_decomposition(const std::vector<std::vector<int>>& G) {
    const int N = G.size();
    std::vector<int> sz(N), level(N, -1), sz_comp(N), par(N);

    auto dfs_size = [&](auto& dfs_size, int v, int p) -> int {
        sz[v] = 1;
        for (int c : G[v]) {
            if (c != p && level[c] == -1) sz[v] += dfs_size(dfs_size, c, v);
        }
        return sz[v];
    };

    auto dfs_centroid = [&](auto& dfs_centroid, int v, int p, int n) -> int {
        for (int c : G[v]) {
            if (c != p && level[c] == -1 && sz[c] > n / 2)
                return dfs_centroid(dfs_centroid, c, v, n);
        }
        return v;
    };

    auto decompose = [&](auto& decompose, int v, int k, int p) -> void {
        int n = dfs_size(dfs_size, v, -1);
        int s = dfs_centroid(dfs_centroid, v, -1, n);
        level[s] = k;
        sz_comp[s] = n;
        par[s] = p;
        for (int c : G[s]) {
            if (level[c] == -1) decompose(decompose, c, k + 1, s);
        }
    };

    decompose(decompose, 0, 0, -1);
    return {level, sz_comp, par};
}
#line 9 "tree/range_contour_aggregation.hpp"

template <typename Group>
class RangeContourAggregation {
    using T = Group::T;

   public:
    RangeContourAggregation() = default;
    explicit RangeContourAggregation(const std::vector<std::vector<int>>& G)
        : seg(G.size()),
          seg_sub(G.size()),
          dist(G.size()),
          idx(G.size()),
          sub(G.size()),
          idx_sub(G.size()),
          first(G.size()),
          first_sub(G.size()) {
        std::tie(level, sz_comp, par) = centroid_decomposition(G);

        for (int v = 0; v < (int)G.size(); ++v) {
            // calculate for each vertex u,
            // dist: dist from v
            // idx: index in the BFS order
            // sub: subtree of v that u belongs to
            // idx_sub: index in the subtree

            // calculate for each depth d,
            // first: first index of depth d
            // first_sub: first index of depth d in the subtree i

            dist[v][v] = 0;
            std::queue<std::pair<int, int>> que;
            que.push({v, -1});
            int k = 0, sub_idx = 0;
            std::vector<int> ks;

            while (!que.empty()) {
                auto [u, i] = que.front();
                que.pop();

                // update component info
                int d = dist[v][u];
                if (d >= (int)first[v].size()) {
                    first[v].push_back(k);
                }
                idx[v][u] = k++;
                // enter subtree
                if (d == 1) {
                    i = sub_idx++;
                    ks.push_back(0);
                    first_sub[v].emplace_back();
                }
                // update subtree info
                if (i != -1) {
                    sub[v][u] = i;
                    if (d - 1 >= (int)first_sub[v][i].size()) {
                        first_sub[v][i].push_back(ks[i]);
                    }
                    idx_sub[v][u] = ks[i]++;
                }

                for (int w : G[u]) {
                    if (level[w] > level[v] && !dist[v].contains(w)) {
                        dist[v][w] = d + 1;
                        que.push({w, i});
                    }
                }
            }

            first[v].push_back(k);
            seg[v] = FenwickTree<Group>(k);
            for (int i = 0; i < sub_idx; ++i) {
                first_sub[v][i].push_back(ks[i]);
                seg_sub[v].emplace_back(ks[i]);
            }
        }
    }

    void update(int v, T x) {
        seg[v].update(0, x);
        for (int p = par[v]; p != -1; p = par[p]) {
            seg[p].update(idx[p][v], x);
            int i = sub[p][v];
            seg_sub[p][i].update(idx_sub[p][v], x);
        }
    }

    T fold(int v, int d) const {
        int dd = std::min((int)first[v].size() - 1, d);
        T res = seg[v].prefix_fold(first[v][dd]);
        for (int p = par[v]; p != -1; p = par[p]) {
            int dep = d - dist[p].at(v);
            if (dep >= 0) {
                dd = std::min((int)first[p].size() - 1, dep);
                res = Group::op(res, seg[p].prefix_fold(first[p][dd]));
                if (dd > 0) {
                    int i = sub[p].at(v);
                    dd = std::min((int)first_sub[p][i].size() - 1, dep - 1);
                    res = Group::op(res, Group::inv(seg_sub[p][i].prefix_fold(
                                             first_sub[p][i][dd])));
                }
            }
        }
        return res;
    }

   private:
    std::vector<int> level, sz_comp, par;
    std::vector<FenwickTree<Group>> seg;
    std::vector<std::vector<FenwickTree<Group>>> seg_sub;
    std::vector<std::unordered_map<int, int>> dist, idx, sub, idx_sub;
    std::vector<std::vector<int>> first;
    std::vector<std::vector<std::vector<int>>> first_sub;
};
Back to top page