Processing math: 100%
My Algorithm : kopricky アルゴリズムライブラリ

kopricky アルゴリズムライブラリ

Path Query / Subtree Query + Ball Query

コードについての説明

木上の頂点の値を (加算 / 更新), パスからの距離が d 以内の頂点の値の (総和 / 最小値), 頂点 u を根とする部分木内の頂点で u からの距離が d 以下であるものの値の (総和 / 最小値) などを処理するアルゴリズム.
yosupo さんのアイデア(参考 URL) を参考にさせていただき 1 層目に weighted balanced な木を用いた.
各クエリ O(log2n) 時間で処理するものの定数が重く, また空間計算量も O(nlogn) であるもののかなり大きくなる. 特に頂点更新 + 範囲最小値の方はオフラインの場合には事前に値を座圧し int 型の値として処理した方が空間計算量的に良い (空間計算量削減のためセグ木の代わりに BIT 2 つで RMQ を処理するテクを用いてみたが実測が遅くなってしまった).

(関数)
vertex_add(u, x) / vertex_update(u, x): 頂点 u の値に x を加算する / 頂点 u の値を x に更新する.
path_query(u, v, d): 頂点 u, v を結ぶパスからの距離が d 以内の頂点についてクエリに答える.
subtree_query(u, d): 頂点 u を根とする部分木内の頂点で u からの距離が d 以下であるものについてクエリに答える.
例えば path_query(u, v, 0) は path query, path_query(u, u, d) は ball query に対応する.

時間計算量: 構築 O(nlogn), 各クエリ O(log2n), 空間計算量 O(nlogn)

コード(頂点加算 + 範囲総和)

  1. template<typename T> class BIT {
  2. private:
  3.   int n;
  4.   vector<T> bit;
  5.     T sm;
  6. public:
  7.     void init(const vector<T>& v){
  8.         n = (int)v.size() + 1, bit.resize(n, 0);
  9.         sm = accumulate(v.begin(), v.end(), (T)0);
  10.         int j;
  11.         for(int i = 1; i < n; ++i){
  12.             bit[i] += v[i - 1], j = i + (i & -i);
  13.             if(j < n) bit[j] += bit[i];
  14.         }
  15.     }
  16.     void add(int i, const T x){
  17.         ++i, sm += x;
  18.         while(i < n) bit[i] += x, i += i & -i;
  19.     }
  20.     T sum(int i) const {
  21.         if(i == -1) return 0;
  22.         else if(i == n - 2) return sm;
  23.         else{
  24.             ++i;
  25.             T s = 0;
  26.             while(i > 0) s += bit[i], i -= i & -i;
  27.             return s;
  28.         }
  29.     }
  30. };
  31.  
  32. template<typename T> class HLdecomposition {
  33. private:
  34.     struct node {
  35.         int left, right, lb, rb;
  36.         BIT<T> bt[3];
  37.         vector<pair<int, int> > pvec[3];
  38.     };
  39.     int V;
  40.     vector<vector<int> > G;
  41.     vector<node> nodes;
  42.     vector<T> value;
  43.     vector<int> sz, depth, pathup, path_id, node_id, wbt_root;
  44.     void predfs(const int u, const int p, const int d){
  45.         sz[u] = 1, depth[u] = d;
  46.         for(int& v : G[u]){
  47.             if(v == p){
  48.                 if(v == G[u].back()) break;
  49.                 else swap(v, G[u].back());
  50.             }
  51.             predfs(v, u, d + 1);
  52.             sz[u] += sz[v];
  53.             if(sz[v] > sz[G[u][0]]) swap(v, G[u][0]);
  54.         }
  55.     }
  56.     void dfs(const int u, const int p, const int par, int& tm, vector<vector<int> >& hpath){
  57.         for(const int v : G[u]){
  58.             if(v == p) continue;
  59.             if(v == G[u][0]){
  60.                 pathup[v] = par, path_id[v] = tm, hpath[tm].push_back(v);
  61.                 dfs(v, u, par, tm, hpath);
  62.             }else{
  63.                 pathup[v] = u, path_id[v] = ++tm, hpath.push_back({v});
  64.                 dfs(v, u, u, tm, hpath);
  65.             }
  66.         }
  67.     }
  68.     void node_merge(BIT<T>& bt, vector<pair<int, int> >& p, const vector<pair<int, int> >& q, const vector<pair<int, int> >& r){
  69.         const int n = (int)q.size() + (int)r.size();
  70.         p.resize(n);
  71.         merge(q.begin(), q.end(), r.begin(), r.end(), p.begin());
  72.         vector<T> val(n);
  73.         for(int i = 0; i < n; ++i) val[i] = value[p[i].second];
  74.         bt.init(val);
  75.     }
  76.     int build_rec(const int l, const int r, const vector<int>& vers, const vector<long long>& sm, const vector<vector<vector<int> > >& info, int& id){
  77.         const int k = id++;
  78.         nodes[k].lb = l, nodes[k].rb = r;
  79.         if(r - l == 1){
  80.             const int u = vers[l];
  81.             for(int i = 0; i < 3; ++i) nodes[k].pvec[i].resize(sz[u]);
  82.             vector<T> val(sz[u]);
  83.             int id = 0;
  84.             for(int i = 0; i < (int)info[u].size(); ++i){
  85.                 for(const int v : info[u][i]){
  86.                     nodes[k].pvec[0][id] = {i - l, v}, nodes[k].pvec[1][id] = {i, v}, nodes[k].pvec[2][id] = {i + l, v};
  87.                     val[id++] = value[v];
  88.                 }
  89.             }
  90.             for(int i = 0; i < 3; ++i) nodes[k].bt[i].init(val);
  91.             return node_id[u] = k;
  92.         }
  93.         int mid = l + 1;
  94.         long long cri = min(sm[mid] - sm[l], sm[r] - sm[mid]);
  95.         for(int i = l + 2; i < r; ++i){
  96.             if(cri < min(sm[i] - sm[l], sm[r] - sm[i])){
  97.                 mid = i, cri = min(sm[mid] - sm[l], sm[r] - sm[mid]);
  98.             }
  99.         }
  100.         const int left = build_rec(l, mid, vers, sm, info, id), right = build_rec(mid, r, vers, sm, info, id);
  101.         nodes[k].left = left, nodes[k].right = right;
  102.         for(int i = 0; i < 3; ++i) node_merge(nodes[k].bt[i], nodes[k].pvec[i], nodes[left].pvec[i], nodes[right].pvec[i]);
  103.         return k;
  104.     }
  105.     int build_wbt(const vector<int>& vers, int& id, const vector<vector<vector<int> > >& info){
  106.         int n = (int)vers.size();
  107.         vector<long long> sm(n + 1, 0);
  108.         for(int i = 0; i < n; ++i) sm[i + 1] = sm[i] + sz[vers[i]];
  109.         return build_rec(0, (int)vers.size(), vers, sm, info, id);
  110.     }
  111.     void build_impl(const int root){
  112.         predfs(root, -1, 0);
  113.         int tm = 0;
  114.         vector<vector<int> > hpath;
  115.         pathup[root] = -1, path_id[root] = 0, hpath.push_back({root});
  116.         dfs(root, -1, -1, tm, hpath);
  117.         vector<vector<vector<int> > > info(V);
  118.         for(int i = 0; i < V; ++i){
  119.             const int d = depth[i];
  120.             for(int u = i; u >= 0; u = pathup[u]){
  121.                 if((int)info[u].size() <= d - depth[u]) info[u].resize(d - depth[u] + 1);
  122.                 info[u][d - depth[u]].push_back(i);
  123.             }
  124.         }
  125.         int id = 0;
  126.         for(const vector<int>& path : hpath){
  127.             for(int i = 0; i < (int)path.size() - 1; ++i) sz[path[i]] -= sz[path[i + 1]];
  128.             wbt_root.push_back(build_wbt(path, id, info));
  129.         }
  130.     }
  131.     void vertex_add_impl(const int u, const T x){
  132.         for(int v = u; v >= 0; v = pathup[v]){
  133.             const int d = depth[u] - depth[v];
  134.             const int id = nodes[node_id[v]].lb;
  135.             int w = wbt_root[path_id[v]];
  136.             while(true){
  137.                 node& cur = nodes[w];
  138.                 cur.bt[0].add((int)(lower_bound(cur.pvec[0].begin(), cur.pvec[0].end(), make_pair(d - id, u)) - cur.pvec[0].begin()), x);
  139.                 cur.bt[1].add((int)(lower_bound(cur.pvec[1].begin(), cur.pvec[1].end(), make_pair(d, u)) - cur.pvec[1].begin()), x);
  140.                 cur.bt[2].add((int)(lower_bound(cur.pvec[2].begin(), cur.pvec[2].end(), make_pair(d + id, u)) - cur.pvec[2].begin()), x);
  141.                 if(nodes[w].rb - nodes[w].lb == 1) break;
  142.                 else if(id < nodes[nodes[w].left].rb) w = nodes[w].left;
  143.                 else w = nodes[w].right;
  144.             }
  145.         }
  146.     }
  147.     T query(int a, int b, const int x, const node& cur, int t) const {
  148.         const int l = cur.lb, r = cur.rb;
  149.         if(r <= a || b <= l) return 0;
  150.         if(a <= l && r <= b){
  151.             return cur.bt[t].sum((int)(upper_bound(cur.pvec[t].begin(), cur.pvec[t].end(), make_pair(x, V)) - cur.pvec[t].begin()) - 1);
  152.         }else{
  153.             return query(a, b, x, nodes[cur.left], t) + query(a, b, x, nodes[cur.right], t);
  154.         }
  155.     }
  156.     T path_query_impl(int u, int v, const int d) const {
  157.         T res = 0;
  158.         while(path_id[u] != path_id[v]){
  159.             if(path_id[u] < path_id[v]){
  160.                 const int id = nodes[node_id[v]].lb;
  161.                 const node& wroot = nodes[wbt_root[path_id[v]]];
  162.                 res += query(0, id + 1, d, wroot, 1) + query(id + 1, wroot.rb, d + id, wroot, 2) - query(0, wroot.rb, d - 1, wroot, 2);
  163.                 v = pathup[v];
  164.             }else{
  165.                 const int id = nodes[node_id[u]].lb;
  166.                 const node& wroot = nodes[wbt_root[path_id[u]]];
  167.                 res += query(0, id + 1, d, wroot, 1) + query(id + 1, wroot.rb, d + id, wroot, 2) - query(0, wroot.rb, d - 1, wroot, 2);
  168.                 u = pathup[u];
  169.             }
  170.         }
  171.         if(depth[u] > depth[v]) swap(u, v);
  172.         const int L = nodes[node_id[u]].lb, R = nodes[node_id[v]].lb;
  173.         const node& wroot = nodes[wbt_root[path_id[u]]];
  174.         res += query(0, L, d - L, wroot, 0) + query(L, R + 1, d, wroot, 1) + query(R + 1, wroot.rb, d + R, wroot, 2);
  175.         const int lca = depth[u];
  176.         int prev_wbt_root = wbt_root[path_id[u]];
  177.         for(u = pathup[u]; u >= 0 && depth[u] >= lca - d; prev_wbt_root = wbt_root[path_id[u]], u = pathup[u]){
  178.             const int id = nodes[node_id[u]].lb;
  179.             const node& wroot = nodes[wbt_root[path_id[u]]];
  180.             res += query(0, id + 1, d - lca + depth[u] - id, wroot, 0) + query(id + 1, wroot.rb, d - lca + depth[u] + id, wroot, 2);
  181.             const node& prev_wroot = nodes[prev_wbt_root];
  182.             res -= query(0, prev_wroot.rb, d - lca + depth[u] - 1, prev_wroot, 2);
  183.         }
  184.         return res;
  185.     }
  186.     T subtree_query_impl(const int u, const int d) const {
  187.         const int id = nodes[node_id[u]].lb;
  188.         const node& wroot = nodes[wbt_root[path_id[u]]];
  189.         return query(id, wroot.rb, d + id, wroot, 2);
  190.     }
  191. public:
  192.     HLdecomposition(const vector<T>& _value)
  193.          : V((int)_value.size()), G(V), nodes(2 * V - 1), value(_value),
  194.              sz(V), depth(V), pathup(V), path_id(V), node_id(V){}
  195.     void add_edge(const int u, const int v){
  196.         G[u].push_back(v), G[v].push_back(u);
  197.     }
  198.     void build(const int root=0){
  199.         build_impl(root);
  200.     }
  201.     void vertex_add(const int u, const T x){
  202.         vertex_add_impl(u, x);
  203.     }
  204.     T path_query(const int u, const int v, const int d) const {
  205.         assert(d >= 0);
  206.         return path_query_impl(u, v, d);
  207.     }
  208.     T subtree_query(const int u, const int d) const {
  209.         assert(d >= 0);
  210.         return subtree_query_impl(u, d);
  211.     }
  212. };

コード(頂点更新 + 範囲最小値)

  1. template<typename T> class segtree {
  2. private:
  3.     int n, sz;
  4.     vector<T> node;
  5. public:
  6.     void init(const vector<T>& v){
  7.         n = 1, sz = (int)v.size();
  8.         while(n < sz) n *= 2;
  9.         node.assign(2 * n, numeric_limits<T>::max());
  10.         for(int i = 0; i < sz; ++i) node[i + n] = v[i];
  11.         for(int i = n - 1; i >= 1; --i) node[i] = min(node[i << 1], node[(i << 1) | 1]);
  12.     }
  13.     void update(int k, const T a){
  14.         const T old = node[k += n];
  15.         if(old == a) return;
  16.         node[k] = a;
  17.         while(k >>= 1){
  18.             if(node[k] > a) node[k] = a;
  19.             else if(node[k] == old) node[k] = min(node[k << 1], node[(k << 1) | 1]);
  20.             else break;
  21.         }
  22.     }
  23.     T query(int a, int b) const {
  24.         if(a == 0 && b == sz) return node[1];
  25.         T res1 = numeric_limits<T>::max(), res2 = numeric_limits<T>::max();
  26.         a += n, b += n;
  27.         while(a != b){
  28.             if(a & 1) res1 = min(res1, node[a++]);
  29.             if(b & 1) res2 = min(node[--b], res2);
  30.             a >>= 1, b >>= 1;
  31.         }
  32.         return min(res1, res2);
  33.     }
  34. };
  35.  
  36. template<typename T> class HLdecomposition {
  37. private:
  38.     struct node {
  39.         int left, right, lb, rb;
  40.         segtree<T> seg[3];
  41.         vector<pair<int, int> > pvec[3];
  42.     };
  43.     int V;
  44.     vector<vector<int> > G;
  45.     vector<node> nodes;
  46.     vector<T> value;
  47.     vector<int> sz, depth, pathup, path_id, node_id, wbt_root;
  48.     void predfs(const int u, const int p, const int d){
  49.         sz[u] = 1, depth[u] = d;
  50.         for(int& v : G[u]){
  51.             if(v == p){
  52.                 if(v == G[u].back()) break;
  53.                 else swap(v, G[u].back());
  54.             }
  55.             predfs(v, u, d + 1);
  56.             sz[u] += sz[v];
  57.             if(sz[v] > sz[G[u][0]]) swap(v, G[u][0]);
  58.         }
  59.     }
  60.     void dfs(const int u, const int p, const int par, int& tm, vector<vector<int> >& hpath){
  61.         for(const int v : G[u]){
  62.             if(v == p) continue;
  63.             if(v == G[u][0]){
  64.                 pathup[v] = par, path_id[v] = tm, hpath[tm].push_back(v);
  65.                 dfs(v, u, par, tm, hpath);
  66.             }else{
  67.                 pathup[v] = u, path_id[v] = ++tm, hpath.push_back({v});
  68.                 dfs(v, u, u, tm, hpath);
  69.             }
  70.         }
  71.     }
  72.     void node_merge(segtree<T>& sg, vector<pair<int, int> >& p, const vector<pair<int, int> >& q, const vector<pair<int, int> >& r){
  73.         const int n = (int)q.size() + (int)r.size();
  74.         p.resize(n);
  75.         merge(q.begin(), q.end(), r.begin(), r.end(), p.begin());
  76.         vector<T> val(n);
  77.         for(int i = 0; i < n; ++i) val[i] = value[p[i].second];
  78.         sg.init(val);
  79.     }
  80.     int build_rec(const int l, const int r, const vector<int>& vers, const vector<long long>& sm, const vector<vector<vector<int> > >& info, int& id){
  81.         const int k = id++;
  82.         nodes[k].lb = l, nodes[k].rb = r;
  83.         if(r - l == 1){
  84.             const int u = vers[l];
  85.             for(int i = 0; i < 3; ++i) nodes[k].pvec[i].resize(sz[u]);
  86.             vector<T> val(sz[u]);
  87.             int id = 0;
  88.             for(int i = 0; i < (int)info[u].size(); ++i){
  89.                 for(const int v : info[u][i]){
  90.                     nodes[k].pvec[0][id] = {i - l, v}, nodes[k].pvec[1][id] = {i, v}, nodes[k].pvec[2][id] = {i + l, v};
  91.                     val[id++] = value[v];
  92.                 }
  93.             }
  94.             for(int i = 0; i < 3; ++i) nodes[k].seg[i].init(val);
  95.             return node_id[u] = k;
  96.         }
  97.         int mid = l + 1;
  98.         long long cri = min(sm[mid] - sm[l], sm[r] - sm[mid]);
  99.         for(int i = l + 2; i < r; ++i){
  100.             if(cri < min(sm[i] - sm[l], sm[r] - sm[i])){
  101.                 mid = i, cri = min(sm[mid] - sm[l], sm[r] - sm[mid]);
  102.             }
  103.         }
  104.         const int left = build_rec(l, mid, vers, sm, info, id), right = build_rec(mid, r, vers, sm, info, id);
  105.         nodes[k].left = left, nodes[k].right = right;
  106.         for(int i = 0; i < 3; ++i) node_merge(nodes[k].seg[i], nodes[k].pvec[i], nodes[left].pvec[i], nodes[right].pvec[i]);
  107.         return k;
  108.     }
  109.     int build_wbt(const vector<int>& vers, int& id, const vector<vector<vector<int> > >& info){
  110.         int n = (int)vers.size();
  111.         vector<long long> sm(n + 1, 0);
  112.         for(int i = 0; i < n; ++i) sm[i + 1] = sm[i] + sz[vers[i]];
  113.         return build_rec(0, (int)vers.size(), vers, sm, info, id);
  114.     }
  115.     void build_impl(const int root){
  116.         predfs(root, -1, 0);
  117.         int tm = 0;
  118.         vector<vector<int> > hpath;
  119.         pathup[root] = -1, path_id[root] = 0, hpath.push_back({root});
  120.         dfs(root, -1, -1, tm, hpath);
  121.         vector<vector<vector<int> > > info(V);
  122.         for(int i = 0; i < V; ++i){
  123.             const int d = depth[i];
  124.             for(int u = i; u >= 0; u = pathup[u]){
  125.                 if((int)info[u].size() <= d - depth[u]) info[u].resize(d - depth[u] + 1);
  126.                 info[u][d - depth[u]].push_back(i);
  127.             }
  128.         }
  129.         int id = 0;
  130.         for(const vector<int>& path : hpath){
  131.             for(int i = 0; i < (int)path.size() - 1; ++i) sz[path[i]] -= sz[path[i + 1]];
  132.             wbt_root.push_back(build_wbt(path, id, info));
  133.         }
  134.     }
  135.     void vertex_update_impl(const int u, const T x){
  136.         for(int v = u; v >= 0; v = pathup[v]){
  137.             const int d = depth[u] - depth[v];
  138.             const int id = nodes[node_id[v]].lb;
  139.             int w = wbt_root[path_id[v]];
  140.             while(true){
  141.                 node& cur = nodes[w];
  142.                 cur.seg[0].update((int)(lower_bound(cur.pvec[0].begin(), cur.pvec[0].end(), make_pair(d - id, u)) - cur.pvec[0].begin()), x);
  143.                 cur.seg[1].update((int)(lower_bound(cur.pvec[1].begin(), cur.pvec[1].end(), make_pair(d, u)) - cur.pvec[1].begin()), x);
  144.                 cur.seg[2].update((int)(lower_bound(cur.pvec[2].begin(), cur.pvec[2].end(), make_pair(d + id, u)) - cur.pvec[2].begin()), x);
  145.                 if(nodes[w].rb - nodes[w].lb == 1) break;
  146.                 else if(id < nodes[nodes[w].left].rb) w = nodes[w].left;
  147.                 else w = nodes[w].right;
  148.             }
  149.         }
  150.     }
  151.     T query(int a, int b, const int x, const node& cur, int t) const {
  152.         const int l = cur.lb, r = cur.rb;
  153.         if(r <= a || b <= l) return numeric_limits<T>::max();
  154.         if(a <= l && r <= b){
  155.             return cur.seg[t].query(0, (int)(upper_bound(cur.pvec[t].begin(), cur.pvec[t].end(), make_pair(x, V)) - cur.pvec[t].begin()));
  156.         }else{
  157.             return min(query(a, b, x, nodes[cur.left], t), query(a, b, x, nodes[cur.right], t));
  158.         }
  159.     }
  160.     T path_query_impl(int u, int v, const int d) const {
  161.         T res = numeric_limits<T>::max();
  162.         while(path_id[u] != path_id[v]){
  163.             if(path_id[u] < path_id[v]){
  164.                 const int id = nodes[node_id[v]].lb;
  165.                 const node& wroot = nodes[wbt_root[path_id[v]]];
  166.                 res = min({res, query(0, id + 1, d, wroot, 1), query(id + 1, wroot.rb, d + id, wroot, 2)});
  167.                 v = pathup[v];
  168.             }else{
  169.                 const int id = nodes[node_id[u]].lb;
  170.                 const node& wroot = nodes[wbt_root[path_id[u]]];
  171.                 res = min({res, query(0, id + 1, d, wroot, 1), query(id + 1, wroot.rb, d + id, wroot, 2)});
  172.                 u = pathup[u];
  173.             }
  174.         }
  175.         if(depth[u] > depth[v]) swap(u, v);
  176.         const int L = nodes[node_id[u]].lb, R = nodes[node_id[v]].lb;
  177.         const node& wroot = nodes[wbt_root[path_id[u]]];
  178.         res = min({res, query(0, L, d - L, wroot, 0), query(L, R + 1, d, wroot, 1), query(R + 1, wroot.rb, d + R, wroot, 2)});
  179.         const int lca = depth[u];
  180.         for(u = pathup[u]; u >= 0 && depth[u] >= lca - d; u = pathup[u]){
  181.             const int id = nodes[node_id[u]].lb;
  182.             const node& wroot = nodes[wbt_root[path_id[u]]];
  183.             res = min({res, query(0, id + 1, d - lca + depth[u] - id, wroot, 0), query(id + 1, wroot.rb, d - lca + depth[u] + id, wroot, 2)});
  184.         }
  185.         return res;
  186.     }
  187.     T subtree_query_impl(const int u, const int d) const {
  188.         const int id = nodes[node_id[u]].lb;
  189.         const node& wroot = nodes[wbt_root[path_id[u]]];
  190.         return query(id, wroot.rb, d + id, wroot, 2);
  191.     }
  192. public:
  193.     HLdecomposition(const vector<T>& _value)
  194.          : V((int)_value.size()), G(V), nodes(2 * V - 1), value(_value),
  195.              sz(V), depth(V), pathup(V), path_id(V), node_id(V){}
  196.     void add_edge(const int u, const int v){
  197.         G[u].push_back(v), G[v].push_back(u);
  198.     }
  199.     void build(const int root=0){
  200.         build_impl(root);
  201.     }
  202.     void vertex_update(const int u, const T x){
  203.         vertex_update_impl(u, x);
  204.     }
  205.     T path_query(const int u, const int v, const int d) const {
  206.         assert(d >= 0);
  207.         return path_query_impl(u, v, d);
  208.     }
  209.     T subtree_query(const int u, const int d) const {
  210.         assert(d >= 0);
  211.         return subtree_query_impl(u, d);
  212.     }
  213. };

verify 用の問題

手元でかなりテストはした.