読者です 読者をやめる 読者になる 読者になる

mayoko’s diary

プロコンとかいろいろ。

yukicoder No.399 動的な領主(その2)

初 HL 分解です。

解法

HL 分解を使います。下の記事が HL 分解については詳しいです。
math314.hateblo.jp

また, 以下で書かれていることは pekempey さんの記事でもまとめられています。
pekempey.hatenablog.com

下のコードは koyumeishi さんのを参考にして書きました(ほとんど一致)
#103178 No.399 動的な領主 - yukicoder

HL 分解を使うと,

  • 分解後の木グラフの深さは log N で抑えられる
  • コンポーネントは一直線になっているので区間として扱える(つまり segTree とかが使える)

という利点があります。つまり, 木におけるパスをいくつかの区間(そしてその区間の数は O(log N) 個に抑えられる)の問題に分けることができます。
これを今回の問題に適用します。今回の問題では,

  • パスにある頂点それぞれの value に 1 を足す
  • パスにある頂点それぞれの value の和を取る

ということができれば良いです。セグメント木を使うところ, lca を求めるところは良いと思うので,

u -> lca -> v へのパスを考えたとき, HL 分解してどう各区間を求めるかというのを考えます。そのために, HL 分解では,

というのをメモしておきます。

v から lca へのパスを考えます(v が lca に向かって登っていくようなイメージ)。
v と lca が同じコンポーネントになるまでは, v と同じコンポーネントにある頂点で最も上にある頂点まで行くので, そのコンポーネントでは区間 [0, (v のそのコンポーネント内での順番)] というのを区間として考えれば良いです。

で, v を登らせるのですが, これは v と同じコンポーネントで親になっているものをたどれば良いです。

v, u を lca と同じコンポーネントまで登らせた後は, [(v のそのコンポーネント内での順番), (u のそのコンポーネント内での順番)] の区間を考えれば良いです。

class Tree {
public:
    Tree(int V, int root) : V(V), root(root) {
        T.resize(V);
        for (int i = 0; i < MAXLOGV; i++) parent[i].resize(V);
        depth.resize(V);
    }
    // uとvをつなぐ
    // lcaを求めることが主目的なので無向グラフとしている
    void unite(int u, int v) {
        T[u].push_back(v);
        T[v].push_back(u);
    }
    // initする
    // コンストラクタだけじゃなくてこれも呼ばないとlcaが求められないぞ
    void init() {
        dfs(root, -1, 0);
        for (int k = 0; k+1 < MAXLOGV; k++) {
            for (int v = 0; v < V; v++) {
                if (parent[k][v] < 0) parent[k+1][v] = -1;
                else parent[k+1][v] = parent[k][parent[k][v]];
            }
        }
    }
    // uとvのlcaを求める
    int lca(int u, int v) const {
        if (depth[u] > depth[v]) swap(u, v);
        for (int k = 0; k < MAXLOGV; k++) {
            if ((depth[v] - depth[u])>>k & 1) {
                v = parent[k][v];
            }
        }
        if (u == v) return u;
        for (int k = MAXLOGV-1; k >= 0; k--) {
            if (parent[k][u] != parent[k][v]) {
                u = parent[k][u];
                v = parent[k][v];
            }
        }
        return parent[0][u];
    }
    // uとvの距離を求める
    // edgeを定義しないといけない時はこれじゃダメ
    int dist(int u, int v) const {
        int p = lca(u, v);
        return (depth[u]-depth[p]) + (depth[v]-depth[p]);
    }
    void dfs(int v, int p, int d) {
        parent[0][v] = p;
        depth[v] = d;
        for (int next : T[v]) {
            if (next != p) dfs(next, v, d+1);
        }
    }
    static const int MAXLOGV = 25;
    // グラフの隣接リスト表現
    vector<vector<int> > T;
    // 頂点の数
    int V;
    // 根ノードの番号
    int root;

    // 親ノード
    vector<int> parent[MAXLOGV];
    // 根からの深さ
    vector<int> depth;
};


// 遅延評価つきセグメント木
// update: [l, r) の区間に値 v を一様に追加する
// query:  [l, r) の区間の和を求める
struct ST {
    vector<ll> seg, lazy;
    int size;

	ST() {}
    ST(int n) {
        size = 1;
        while (size < n) size *= 2;
        seg.resize(size * 2);
        lazy.resize(size * 2);
    }

    inline void push(int k, int l, int r) {
        if (lazy[k] != 0) {
            seg[k] += lazy[k] * (r - l);
            if (r - l > 1) {
                lazy[k * 2 + 1] += lazy[k];
                lazy[k * 2 + 2] += lazy[k];
            }
            lazy[k] = 0;
        }
    }

    inline ll merge(ll x, ll y) {
        return x + y;
    }

    void update(int a, int b, ll v, int k, int l, int r) {
        push(k, l, r);
        if (r <= a || b <= l) return;
        if (a <= l && r <= b) {
            lazy[k] += v;
            push(k, l, r);
        } else {
            update(a, b, v, k * 2 + 1, l, (l + r) / 2);
            update(a, b, v, k * 2 + 2, (l + r) / 2, r);
            seg[k] = merge(seg[k * 2 + 1], seg[k * 2 + 2]);
        }
    }

    void update(int a, int b, ll v) {
        return update(a, b, v, 0, 0, size);
    }

    ll query(int a, int b, int k, int l, int r) {
        push(k, l, r);
        if (r <= a || b <= l) return 0;
        if (a <= l && r <= b) return seg[k];
        ll x = query(a, b, k * 2 + 1, l, (l + r) / 2);
        ll y = query(a, b, k * 2 + 2, (l + r) / 2, r);
        return merge(x, y);
    }

    ll query(int a, int b) {
        return query(a, b, 0, 0, size);
    }
};


class HeavyLightDecomposition {
public:
	struct heavy_set {
		vector<int> element;
		int depth;
		int parent_vertex;
		heavy_set(int v, int d, int par) : element(1, v), depth(d), parent_vertex(par) {}
	};
	vector<vector<int>> G;
	vector<heavy_set> S;
	vector<int> subtree_size;
	vector<int> set_index;
	vector<int> ele_index;
private:
	int get_subtree_size(int v, int p) {
		int& sz = subtree_size[v];
		if (sz > 0) return sz;
		sz = 1;
		for (int ch : G[v]) if (ch != p) {
			sz += get_subtree_size(ch, v);
		}
		return sz;
	}

	void make_path(int v, int p, int set_id) {
		set_index[v] = set_id;
		ele_index[v] = S[set_id].element.size() - 1;

		int largest_child = -1, maxi = 0;

		for (int ch : G[v]) if (ch != p) {
			if (maxi < get_subtree_size(ch, v)) {
				maxi = subtree_size[ch];
				largest_child = ch;
			}
		}

		for (int ch : G[v]) if (ch != p) {
			if (largest_child == ch) {
				S[set_id].element.push_back(ch);
				make_path(ch, v, set_id);
			}
			else {
				S.emplace_back(ch, S[set_id].depth + 1, v);
				make_path(ch, v, S.size() - 1);
			}
		}
	}

	void init(int root) {
		subtree_size.resize(G.size());
		set_index.resize(G.size());
		ele_index.resize(G.size());

		S.emplace_back(root, 0, root);
		make_path(root, root, 0);

		subtree_size.clear();
	}

public:
	HeavyLightDecomposition(const vector<vector<int>>& G, int root = 0) : G(G) {
		init(root);
	}
	pair<int, int> get_position(int v) {
		return pair<int, int>(set_index[v], ele_index[v]);
	}
};

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
	int N;
	cin >> N;
	vector<vector<int> > G(N);
	Tree tree(N, 0);
	for (int i = 0; i < N - 1; i++) {
		int u, v;
		cin >> u >> v;
		u--; v--;
		G[u].push_back(v);
		G[v].push_back(u);
		tree.unite(u, v);
	}
	tree.init();
	HeavyLightDecomposition hld(G);
	vector<ST> seg;
	for (int i = 0; i < hld.S.size(); i++) {
		seg.emplace_back(hld.S[i].element.size());
	}

	ll ans = 0;
	int Q;
	cin >> Q;
	while (Q--) {
		int u, v;
		cin >> u >> v;
		u--; v--;
		int lca = tree.lca(u, v);
		while (hld.get_position(v).first != hld.get_position(lca).first) {
			pii tmp = hld.get_position(v);
			seg[tmp.first].update(0, tmp.second+1, 1);
			ans += seg[tmp.first].query(0, tmp.second + 1);

			v = hld.S[tmp.first].parent_vertex;
		}
		while (hld.get_position(u).first != hld.get_position(lca).first) {
			pii tmp = hld.get_position(u);
			seg[tmp.first].update(0, tmp.second+1, 1);
			ans += seg[tmp.first].query(0, tmp.second + 1);

			u = hld.S[tmp.first].parent_vertex;
		}
		pii tu = hld.get_position(u), tv = hld.get_position(v);
		int mini = min(tu.second, tv.second), maxi = max(tu.second, tv.second);
		seg[tu.first].update(mini, maxi + 1, 1);
		ans += seg[tu.first].query(mini, maxi + 1);
	}

	cout << ans << endl;

    return 0;
}