AOJ 2667 Tree

オイラーツアー第4弾！！！
でもせっかくなので色んな方法で解いてみた！！！

• Euler ツアー
• クエリの平方分割
• HL 分解
• 重心分解 (未、todo)

問題へのリンク

問題概要

頂点 0 を根とする頂点数 N の根付き木において以下の Q 個のクエリに答えよ。なお初期状態ではすべての辺の重みは 0 である。

• type 0: 二頂点 u-v 間のパスの長さを答える
• type 1: 頂点 u を根とする部分木に含まれるすべての辺のコストを x 増やす (x >= 0)

制約

• 1 <= N, Q <= 150000

解法 1: オイラーツアー

公式解説と同じもの。
もうさすがに見慣れた「根付き木に対するクエリ」はオイラーツアーをとると「区間に対するクエリ」に変換されるというやつなん。

• ツリーの一辺に x を足す
• ツリーの二頂点間のパスの長さを求める

だったら超典型問題だけど、今回はツリーの頂点 v の根付き木の辺すべてに x を足す必要がある。一辺に x を足すのは

• セグツリー (か BIT) の「行きの辺」に +x、「帰りの辺」に -x するだけ (2 点のみ更新)

だが、今回は

• 区間に含まれる「行きの辺すべて」に +x、「帰りの辺すべて」に -x する

必要がある。僕は遅延評価セグメントツリーで実現することにした。各ノードには「その区間の (行きの辺の個数) - (帰りの辺の個数)」の情報も持たせる。

• 更新: 遅延値に x を足す
• 遅延評価: ノードに、遅延値 * 「その区間の (行きの辺の個数) - (帰りの辺の個数)」を足す

という感じの遅延評価で OK。計算量は O(N + QlogN)。

//
// Euler Tour
//
// verified:
//   AOJ 2667 Tree
//     http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2667
//

#include <iostream>
#include <functional>
#include <vector>
#include <queue>
#include <stack>
using namespace std;


// Sparse Table
template<class MeetSemiLattice> struct SparseTable {
    vector<vector<pair<MeetSemiLattice,int> > > dat;
    vector<int> height;
    
    SparseTable() { }
    SparseTable(const vector<MeetSemiLattice> &vec) { init(vec); }
    void init(const vector<MeetSemiLattice> &vec) {
        int n = (int)vec.size(), h = 0;
        while ((1<<h) < n) ++h;
        dat.assign(h, vector<pair<MeetSemiLattice,int> >(1<<h));
        height.assign(n+1, 0);
        for (int i = 2; i <= n; i++) height[i] = height[i>>1]+1;
        for (int i = 0; i < n; ++i) dat[0][i] = {vec[i], i};
        for (int i = 1; i < h; ++i)
            for (int j = 0; j < n; ++j)
                dat[i][j] = min(dat[i-1][j], dat[i-1][min(j+(1<<(i-1)),n-1)]);
    }
    
    pair<MeetSemiLattice,int> get(int a, int b) {
        return min(dat[height[b-a]][a], dat[height[b-a]][b-(1<<height[b-a])]);
    }
};

// Segment Tree
template<class Monoid, class Action> struct SegTree {
    using FuncMonoid = function< Monoid(Monoid, Monoid) >;
    using FuncAction = function< void(Monoid&, Action) >;
    using FuncLazy = function< void(Action&, Action) >;
    FuncMonoid FM;
    FuncAction FA;
    FuncLazy FL;
    Monoid UNITY_MONOID;
    Action UNITY_LAZY;
    int SIZE, HEIGHT;
    vector<Monoid> dat;
    vector<Action> lazy;
    
    SegTree() { }
    SegTree(int n, const FuncMonoid fm, const FuncAction fa, const FuncLazy fl,
            const Monoid &unity_monoid, const Action &unity_lazy)
    : FM(fm), FA(fa), FL(fl), UNITY_MONOID(unity_monoid), UNITY_LAZY(unity_lazy) {
        SIZE = 1; HEIGHT = 0;
        while (SIZE < n) SIZE <<= 1, ++HEIGHT;
        dat.assign(SIZE * 2, UNITY_MONOID);
        lazy.assign(SIZE * 2, UNITY_LAZY);
    }
    void init(int n, const FuncMonoid fm, const FuncAction fa, const FuncLazy fl,
              const Monoid &unity_monoid, const Action &unity_lazy) {
        FM = fm; FA = fa; FL = fl;
        UNITY_MONOID = unity_monoid; UNITY_LAZY = unity_lazy;
        SIZE = 1; HEIGHT = 0;
        while (SIZE < n) SIZE <<= 1, ++HEIGHT;
        dat.assign(SIZE * 2, UNITY_MONOID);
        lazy.assign(SIZE * 2, UNITY_LAZY);
    }
    
    /* set, a is 0-indexed */
    void set(int a, const Monoid &v) { dat[a + SIZE] = v; }
    void build() {
        for (int k = SIZE - 1; k > 0; --k)
            dat[k] = FM(dat[k*2], dat[k*2+1]);
    }
    
    /* update [a, b) */
    inline void evaluate(int k) {
        if (lazy[k] == UNITY_LAZY) return;
        if (k < SIZE) FL(lazy[k*2], lazy[k]), FL(lazy[k*2+1], lazy[k]);
        FA(dat[k], lazy[k]);
        lazy[k] = UNITY_LAZY;
    }
    inline void update(int a, int b, const Action &v, int k, int l, int r) {
        evaluate(k);
        if (a <= l && r <= b)  FL(lazy[k], v), evaluate(k);
        else if (a < r && l < b) {
            update(a, b, v, k*2, l, (l+r)>>1), update(a, b, v, k*2+1, (l+r)>>1, r);
            dat[k] = FM(dat[k*2], dat[k*2+1]);
        }
    }
    inline void update(int a, int b, const Action &v) { update(a, b, v, 1, 0, SIZE); }
    
    /* get [a, b) */
    inline Monoid get(int a, int b, int k, int l, int r) {
        evaluate(k);
        if (a <= l && r <= b)
            return dat[k];
        else if (a < r && l < b)
            return FM(get(a, b, k*2, l, (l+r)>>1), get(a, b, k*2+1, (l+r)>>1, r));
        else
            return UNITY_MONOID;
    }
    inline Monoid get(int a, int b) { return get(a, b, 1, 0, SIZE); }
    inline Monoid operator [] (int a) { return get(a, a+1); }
    
    /* debug */
    void print() {
        for (int i = 0; i < SIZE; ++i) { cout << (*this)[i]; if (i != SIZE) cout << ","; }
        cout << endl;
    }
};

// Euler Tour
using Graph = vector<vector<int> >;
using Node = pair<long long, int>;
const auto fm = [](Node a, Node b) { return Node(a.first + b.first, a.second + b.second); };
const auto fa = [](Node &a, long long d) { a.first += d * a.second; };
const auto fl = [](long long &d, long long e) { d += e; };

struct EulerTour {
    // main results
    Graph tree;
    vector<int> depth;
    vector<int> node; // the node-number of i-th element of Euler-tour
    vector<int> vf, ve; // the index of Euler-tour of node v
    vector<int> eid; // the index of edge e (i*2 + (0: dir to leaf, 1: dir to root))
    
    // sub results
    SparseTable<int> st; // depth (to find LCA)
    
    // segtree
    SegTree<Node, long long> seg;
    
    // initialization
    EulerTour(const Graph &tree_) { init(tree_); }
    void init(const Graph &tree_) {
        tree = tree_;
        int V = (int)tree.size();
        depth.resize(V*2-1); node.resize(V*2-1); vf.resize(V); ve.resize(V); eid.resize((V-1)*2);
        seg.init((V-1)*2, fm, fa, fl, Node(0, 0), 0);
        int k = 0;
        dfs(0, -1, 0, k);
        st.init(depth);
        seg.build();
    }
    
    void dfs(int v, int par, int dep, int &ord) {
        node[ord] = v;
        depth[ord] = dep;
        vf[v] = ve[v] = ord;
        ++ord;
        for (auto e : tree[v]) {
            if (e == par) continue;
            seg.set(ord-1, Node(0, 1));
            dfs(e, v, dep+1, ord);
            node[ord] = v;
            depth[ord] = dep;
            ve[v] = ord;
            seg.set(ord-1, Node(0, -1));
            ++ord;
        }
    }
    
    inline int LCA(int u, int v) {
        int a = vf[u], b = vf[v];
        if (a > b) swap(a, b);
        return node[st.get(a, b+1).second];
    }
    
    inline void update(int v, long long x) {
        seg.update(vf[v], ve[v], x);
    }
    
    inline long long get(int v) {
        return seg.get(0, vf[v]).first;
    }
    
    inline long long get(int u, int v) {
        int lca = LCA(u, v);
        return get(u) + get(v) - get(lca)*2;
    }
};

int main () {
    int N, Q;
    cin >> N >> Q;
    Graph G(N);
    for (int i = 0; i < N-1; ++i) {
        int a, b; scanf("%d %d", &a, &b);
        G[a].push_back(b);
        G[b].push_back(a);
    }

    EulerTour et(G);
    
    for (int q = 0; q < Q; ++q) {
        int type; cin >> type;
        if (type == 0) {
            int u, v; scanf("%d %d", &u, &v);
            long long res = et.get(u, v);
            printf("%lld\n", res);
        }
        else {
            int u, x; scanf("%d %d", &u, &x);
            et.update(u, x);
        }
    }
}

解法 2: クエリの平方分割

しょラーさんの記事を参考に。アイディアは

• クエリを √Q 個のブロックに分ける (各ブロックに √Q 個のクエリ)
• 各ブロックだけをみると、そこに関係する頂点数は O(√Q) 個のみ
• 関係する頂点だけを抜き取った森を作り、その森上で愚直なクエリ処理を行う (多くの愚直処理は O(√Q) くらいの計算量でできる)
• 全体の計算量は、各ブロックについての森作りが O(√QN)、クエリ処理が全体で O(Q√Q) くらいになる。

#include <iostream>
#include <vector>
using namespace std;
typedef vector<vector<int> > Graph;

// LCA
struct LCA {
  vector<vector<int> > parent; // parent[d][v] := 2^d-th parent of v
  vector<int> depth;
  LCA() { }
  LCA(const Graph &G, int r = 0) { init(G, r); }
  void init(const Graph &G, int r = 0) {
    int V = (int)G.size();
    int h = 1;
    while ((1<<h) < V) ++h;
    parent.assign(h, vector<int>(V, -1));
    depth.assign(V, -1);
    dfs(G, r, -1, 0);
    for (int i = 0; i+1 < (int)parent.size(); ++i) 
      for (int v = 0; v < V; ++v) 
        if (parent[i][v] != -1) 
          parent[i+1][v] = parent[i][parent[i][v]];
  }  
  void dfs(const Graph &G, int v, int p, int d) {
    parent[0][v] = p;
    depth[v] = d;
    for (auto e : G[v]) if (e != p) dfs(G, e, v, d+1);
  } 
  int get(int u, int v) {
    if (depth[u] > depth[v]) swap(u, v);
    for (int i = 0; i < (int)parent.size(); ++i) 
      if ( (depth[v] - depth[u]) & (1<<i) ) 
        v = parent[i][v];
    if (u == v) return u;
    for (int i = (int)parent.size()-1; i >= 0; --i) {
      if (parent[i][u] != parent[i][v]) {
        u = parent[i][u];
        v = parent[i][v];
      }
    }
    return parent[0][u];
  }  
};

// クエリの平方分割などにともなうグラフの縮約
// use[v] := 1(v を使う)、0(v を使わない)
struct ContractGraph {
  /* common */
  Graph G;           // original graph
  vector<int> depth; // depth of v in original graph

  /* contracted graph */
  vector<int> use;    // 1: use, 0: not use
  vector<int> parent; // parent of contracted graph

  /* additional member */
  vector<long long> fixedWeights; // 各ブロック終了時におけるその頂点までの重み

  /* constructor */
  ContractGraph(const Graph &G_, const vector<int> &depth_) : G(G_), depth(depth_) {
    fixedWeights.assign((int)G.size(), 0);
  }

  /* build contracted graph */
  void build(const vector<int> &use_, int root = 0) {
    use = use_;
    int V = (int)G.size();
    parent.assign(V, -1);
    dfs(root, -1, -1);
  }
  void dfs(int v, int par, int prev_use) {
    int cur_use = prev_use;
    if (use[v]) parent[v] = prev_use, cur_use = v;
    for (auto e : G[v]) if (e != par) dfs(e, v, cur_use);
  }

  /* each process for each problem */
  void propagate(const vector<long long> &add, int v = 0, int par = -1, long long sum = 0, long long sumsum = 0) {
    fixedWeights[v] += sumsum;
    for (auto e : G[v]) {
      if (e == par) continue;
      long long nextsum = sum + add[v];
      long long nextsumsum = sumsum + nextsum;
      propagate(add, e, v, nextsum, nextsumsum);
    }
  }

  long long get(int v, const vector<long long> &add) {
    long long res = fixedWeights[v];
    for (int cur = v; cur != -1; cur = parent[cur]) {
      res += add[cur] * (depth[v] - depth[cur]);
    }
    return res;
  }
};

int main() {
  int N, Q;
  cin >> N >> Q;
  Graph G(N);
  for (int i = 0; i < N-1; ++i) {
    int a, b; scanf("%d %d", &a, &b);
    G[a].push_back(b); G[b].push_back(a);
  }
  LCA lca(G); // lca

  const int sQ = 1000;
  vector<int> use(N, 0);
  ContractGraph c(G, lca.depth);
  vector<int> tempType, tempU, tempV; // ブロック分
  for (int q = 0; q < Q; ++q) {
    int type, u, v; scanf("%d %d %d", &type, &u, &v);
    tempType.push_back(type); tempU.push_back(u); tempV.push_back(v);
    if (type == 0) {
      int p = lca.get(u, v);
      use[u] = use[v] = use[p] = 1;
    }
    else use[u] = 1;

    if ((q + 1) % sQ == 0 || q == Q-1) {
      c.build(use);
      vector<long long> add(N, 0);
      for (int i = 0; i < (int)tempType.size(); ++i) {
        int type = tempType[i];
        if (type == 0) {
          int u = tempU[i], v = tempV[i], p = lca.get(u, v);
          long long res = c.get(u, add) + c.get(v, add) - c.get(p, add)*2;
          printf("%lld\n", res);
        }
        else add[tempU[i]] += tempV[i];
      }
      use.assign(N, 0);
      c.propagate(add);
      tempType.clear(); tempU.clear(); tempV.clear();
    }
  }
}

解法 3: HL 分解

HL分解は、ツリーを

• 高々 O(logV) 本のパスに分解する
• このときツリーの元の edge は以下の 2 種にわかれる。
  • heavy: 分解された各パス上の edge
  • light: パスとパスの先頭同士をつなぐ edge
• 各パスごとにセグ木なりを載せて処理する、light edge は高々 log オーダー個しかないので個別に処理する

という感じのことをする。おおむね 1 クエリあたり の計算量になる。

HL 分解を頑張る流れは、パスに関する問題であれば

• まずは普通の列に対するクエリ処理方法を考える (今回は「区間加算」と「区間和取得」ができればよい)
• HL 分解ライブラリさえあれば、列に対する処理をほぼそのままラムダ式で渡せばいい感じになる

という感じにできるので、実質「列に対するクエリ処理方法を考える」問題に落ちる。Euler ツアーと比較すると「逆演算」を考えなくてよいので汎用性がさらに高そう。

しかし今回は「v を根とする部分木全体に含まれる edge 全体に加算する」というクエリなので一瞬悩む。が、実は単純で以下のようにすればよい (vid, subsize の定義は実装参照)。

• vid[v] と vid[v] * subsize[v] との間の区間に一様加算する

こうして我々は、パスに対するクエリだけでなく、部分木に対するクエリも HL 分解で処理する手段を獲得することができた。

#include <iostream>
#include <functional>
#include <vector>
#include <queue>
#include <stack>
using namespace std;

// BIT
template <class Abel> struct BIT {
    vector<Abel> dat[2];
    Abel UNITY_SUM = 0;                        // to be set
    
    /* [1, n] */
    BIT(int n) { init(n); }
    void init(int n) { for (int iter = 0; iter < 2; ++iter) dat[iter].assign(n + 1, UNITY_SUM); }

    /* a is 1-indexed */
    inline void sub_add(int p, int a, Abel x) {
        for (int i = a; i < (int)dat[p].size(); i += i & -i)
            dat[p][i] = dat[p][i] + x;
    }
    inline void add(int a, int b, Abel x) {
        sub_add(0, a, x * -(a - 1)); sub_add(1, a, x); sub_add(0, b, x * (b - 1)); sub_add(1, b, x * (-1));
    }
    
    /* [1, a], a is 1-indexed */
    inline Abel sub_sum(int p, int a) {
        Abel res = UNITY_SUM;
        for (int i = a; i > 0; i -= i & -i) res = res + dat[p][i];
        return res;
    }
    inline Abel sum(int a, int b) {
        return sub_sum(0, b - 1) + sub_sum(1, b - 1) * (b - 1) - sub_sum(0, a - 1) - sub_sum(1, a - 1) * (a - 1);
    }
    
    /* debug */
    void print() {
        for (int i = 1; i < (int)dat[0].size(); ++i) cout << sum(i, i + 1) << ",";
        cout << endl;
    }
};

// HL-Decomposition
// vid: id of v after HL-Decomposition
// inv: inv[vid[v]] = v
// par: id of parent
// depth
// subsize: size of subtree
// head: head-id in the heavy-path
// prev, next: prev-id, next-id in the heavy-path
// type: the id of tree for forest
// vend: the last-id of node in v-subtree
typedef vector<vector<int> > Graph;
struct HLDecomposition {
    int n;
    Graph G;
    vector<int> vid, inv, par, depth, subsize, head, prev, next, type;
    
    // construct
    HLDecomposition() { }
    HLDecomposition(const Graph &G_) :
        n((int)G_.size()), G(G_),
        vid(n, -1), inv(n), par(n), depth(n), subsize(n, 1),
        head(n), prev(n, -1), next(n, -1), type(n) { }
    void build(vector<int> roots = {0}) {
        int curtype = 0, pos = 0;
        for (auto r : roots) decide_heavy_edge(r), reconstruct(r, curtype++, pos);
    }
    void decide_heavy_edge(int r) {
        stack<pair<int,int> > st;
        par[r] = -1, depth[r] = 0;
        st.emplace(r, 0);
        while (!st.empty()) {
            int v = st.top().first;
            int &i = st.top().second;
            if (i < (int)G[v].size()) {
                int e = G[v][i++];
                if (e == par[v]) continue;
                par[e] = v, depth[e] = depth[v] + 1;
                st.emplace(e, 0);
            }
            else {
                st.pop();
                int maxsize = 0;
                for (auto e : G[v]) {
                    if (e == par[v]) continue;
                    subsize[v] += subsize[e];
                    if (maxsize < subsize[e]) maxsize = subsize[e], prev[e] = v, next[v] = e;
                }
            }
        }
    }
    void reconstruct(int r, int curtype, int &pos) {
        stack<int> st({r});
        while (!st.empty()) {
            int start = st.top(); st.pop();
            for (int v = start; v != -1; v = next[v]) {
                type[v] = curtype;
                vid[v] = pos++;
                inv[vid[v]] = v;
                head[v] = start;
                for (auto e : G[v]) if (e != par[v] && e != next[v]) st.push(e);
            }
        }
    }
    
    // node query [u, v], f([left, right])
    void foreach_nodes(int u, int v, const function<void(int,int)> &f) {
        while (true) {
            if (vid[u] > vid[v]) swap(u, v);
            f(max(vid[head[v]], vid[u]), vid[v]);
            if (head[u] != head[v]) v = par[head[v]];
            else break;
        }
    }
    
    // edge query [u, v], f([left, right])
    void foreach_edges(int u, int v, const function<void(int,int)> &f) {
        while (true) {
            if (vid[u] > vid[v]) swap(u, v);
            if (head[u] != head[v]) {
                f(vid[head[v]], vid[v]);
                v = par[head[v]];
            }
            else {
                if (u != v) {
                    f(vid[u]+1, vid[v]);
                }
                break;
            }
        }
    }

    // LCA
    int lca(int u, int v) {
        while (true) {
            if (vid[u] > vid[v]) swap(u, v);
            if (head[u] == head[v]) return u;
            v = par[head[v]];
        }
    }
};

int main() {
    int N, Q; cin >> N >> Q;
    Graph G(N);
    for (int i = 0; i < N-1; ++i) {
        int u, v; scanf("%d %d", &u, &v);
        G[u].push_back(v); G[v].push_back(u);
    }
    HLDecomposition hld(G);
    hld.build();
    BIT<long long> bit(N+1);
    for (int q = 0; q < Q; ++q) {
        int type, a, b; scanf("%d %d %d", &type, &a, &b);
        if (type == 0) {
            long long res = 0;
            hld.foreach_edges(a, b, [&](int l, int r){ res += bit.sum(l, r+1); });
            cout << res << endl;
        }
        else {
            bit.add(hld.vid[a]+1, hld.vid[a]+hld.subsize[a], b);
        }
    }
}

解法 4: 重心分解

重心分解は「ツリーを再帰的に重心で分解していくと再帰深さが log オーダーになる」というものなので、言わば、ツリーを平衡なものにするような感じになる。

todo