結果

問題 No.650 行列木クエリ
ユーザー simkarensimkaren
提出日時 2020-03-31 09:59:52
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 525 ms / 2,000 ms
コード長 7,097 bytes
コンパイル時間 3,011 ms
コンパイル使用メモリ 211,408 KB
実行使用メモリ 71,088 KB
最終ジャッジ日時 2023-09-06 05:33:29
合計ジャッジ時間 6,064 ms
ジャッジサーバーID
（参考情報） 		judge15 / judge13
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
4,380 KB
testcase_01 AC 149 ms
17,028 KB
testcase_02 AC 508 ms
71,088 KB
testcase_03 AC 2 ms
4,376 KB
testcase_04 AC 151 ms
16,932 KB
testcase_05 AC 525 ms
71,088 KB
testcase_06 AC 1 ms
4,380 KB
testcase_07 AC 2 ms
4,380 KB
testcase_08 AC 128 ms
16,992 KB
testcase_09 AC 307 ms
68,576 KB
testcase_10 AC 1 ms
4,376 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>

using namespace std;

#define ll long long

// #define TEST

const ll mod = 1000000007LL;

struct Matrix {
	ll e00, e01, e10, e11;
	Matrix() { e00 = e11 = 1; e01 = e10 = 0; }
	Matrix(ll e00, ll e01, ll e10, ll e11) : e00(e00), e01(e01), e10(e10), e11(e11) {}
	Matrix operator *(Matrix a) {
		return Matrix(
			(e00 * a.e00 % mod + e01 * a.e10 % mod) % mod,
			(e00 * a.e01 % mod + e01 * a.e11 % mod) % mod,
			(e10 * a.e00 % mod + e11 * a.e10 % mod) % mod,
			(e10 * a.e01 % mod + e11 * a.e11 % mod) % mod
		);
	}
};

template<typename T>
struct SegmentTree {
private:
	int sz, n;
	vector<T> data;
	function<T(T, T)> f;
	T identity_element;
public:
	/* constructor */
	SegmentTree(
		vector<T>& v, // initial data
		T identity_element, // identity element
		function<T(T, T)> f // operation
	) {
		sz = v.size();
		n = 1; while (n < sz) n <<= 1;
		this->f = f;
		this->identity_element = identity_element;
		data.resize(2 * n - 1, identity_element);
		for (int i = 0; i < sz; ++i)
			data[i + n - 1] = v[i];
		for (int i = n - 2; i >= 0; --i)
			data[i] = f(data[2 * i + 1], data[2 * i + 2]);
	}
	/* update query */
	void update(int idx, const T val) {
		idx += n - 1;
		data[idx] = val;
		while (idx) {
			idx = (idx - 1) >> 1;
			data[idx] = f(data[2 * idx + 1], data[2 * idx + 2]);
		}
	}
	/* get query */
	T get(int left, int right, int k = 0, int l = 0, int r = -1) {
		if (r < 0) r = n;
		if (r <= left || right <= l) return identity_element;
		if (left <= l && r <= right) return data[k];
		T val_l = get(left, right, 2 * k + 1, l, (l + r) / 2);
		T val_r = get(left, right, 2 * k + 2, (l + r) / 2, r);
		return f(val_l, val_r);
	}
};

template <typename T>
struct HLD {
	vector<vector<pair<int, T>>> g;
	int n;
	vector<int> subsz, depth, parent;
	vector<vector<T>> chain_edge;
	vector<vector<int>> chain_node;
	map<int, int> chain_top, chain_id, id_in_chain;

	/* constructor */
	HLD(vector<vector<pair<int, T>>>& g) {
		this->g = g; n = g.size();
		parent.resize(g.size());
		subsz.resize(g.size());
		depth.resize(g.size());
	}

	/* init parent, depth, and subsz */
	int dfs_sz(int s = 0, int p = -1, int dep = 0) {
		parent[s] = p;
		depth[s] = dep;
		int sz = 1;
		if (g[s].size() > 2 && g[s][0].first == p)
			swap(g[s][0], g[s][1]); // g[s][0] must be heavy
		for (int i = 0; i < g[s].size(); ++i) {
			if (g[s][i].first != p) {
				sz += dfs_sz(g[s][i].first, s, dep + 1);
				if (subsz[g[s][0].first] < subsz[g[s][i].first])
					swap(g[s][0], g[s][i]); // g[s][0] must be heavy
			}
		}
		subsz[s] = sz;
		return sz;
	}

	/* HL Decomposition */
	void dfs_hld(int s, int& cid) {
		chain_id[s] = cid;
		for (int i = 0; i < g[s].size(); ++i) {
			if (g[s][i].first != parent[s]) {
				if (i == 0) {
					chain_edge.resize(cid + 1);
					chain_node.resize(cid + 1);
					chain_edge[cid].push_back(g[s][i].second);
					chain_node[cid].push_back(g[s][i].first);
					id_in_chain[g[s][i].first] = chain_edge[cid].size();
					dfs_hld(g[s][i].first, cid);
				}
				else {
					++cid;
					chain_edge.resize(cid + 1);
					chain_node.resize(cid + 1);
					chain_edge[cid].push_back(g[s][i].second);
					chain_node[cid].push_back(g[s][i].first);
					id_in_chain[g[s][i].first] = chain_edge[cid].size();
					dfs_hld(g[s][i].first, cid);
				}
			}
		}
	}

	/* the number of chain */
	int chain_number() {
		return chain_node.size();
	}

	/* find chain_top */
	void find_chain_top() {
		for (int i = 0; i < chain_number(); ++i)
			chain_top[i] = parent[chain_node[i][0]];
	}

	/* for test */
	void show() {
#ifdef TEST
		for (int i = 0; i < chain_number(); ++i) {
			cout << "chain id : " << i << endl;
			cout << "    " << chain_top[i];
			for (int j = 0; j < chain_edge[i].size(); ++j)
				cout << "-" << chain_node[i][j];
			cout << endl;
		}
		for (int i = 0; i < n; ++i) {
			cout << "node : " << i << endl;
			cout << "    chain : " << chain_id[i] << ",  id_in_chain : " << id_in_chain[i] << endl;
		}
#endif
	}

	vector<SegmentTree<Matrix>> segs;
	vector<SegmentTree<Matrix>> rev_segs;


	/* build */
	void build() {
		dfs_sz();
		int cid = 0;
		dfs_hld(0, cid);
		find_chain_top();
		/* 【TODO】*/
		// セグ木の準備等
		for (int i = 0; i < chain_number(); ++i) {
			segs.push_back(
				SegmentTree<Matrix>(
					chain_edge[i], Matrix(), [](Matrix a, Matrix b) {return b * a; }
					)
			);
			rev_segs.push_back(
				SegmentTree<Matrix>(
					chain_edge[i], Matrix(), [](Matrix a, Matrix b) {return a * b; }
					)
			);
		}
	}

	/* find lca */
	int LCA(int s, int t) {
		if (chain_id[s] < chain_id[t]) swap(s, t);
		while (chain_id[s] != chain_id[t]) {
			s = chain_top[chain_id[s]];
			if (chain_id[s] < chain_id[t]) swap(s, t);
		}
		if (depth[s] < depth[t]) return s;
		return t;
	}

	/* sub query */
	T query_in_chain(int cid, int s, int t) {
		// s <= t, [s, t]
		return segs[cid].get(s, t + 1);
	}
	T query_in_chain_reverse(int cid, int s, int t) {
		// t <= s, [t, s]
		// return query_in_chain(int cid, int t, int s);
		return rev_segs[cid].get(t, s + 1);
	}
	void chain_update(int cid, int id, T d) {
		segs[cid].update(id, d);
		rev_segs[cid].update(id, d);
	}
	T merge(T l, T r) {
		return r * l;
	}

	/* edge query */
	T edge_query(int s, int t, T ei) {
		int lca = LCA(s, t);
		// s to lca
		T left = ei;
		while (s != lca) {
			if (chain_id[s] != chain_id[lca]) {
				T tmp = query_in_chain_reverse(
					chain_id[s], id_in_chain[s] - 1, 0
				);
				left = merge(left, tmp);
				s = chain_top[chain_id[s]];
			}
			else {
				T tmp = query_in_chain_reverse(
					chain_id[s], id_in_chain[s] - 1, id_in_chain[lca]
				);
				left = merge(left, tmp);
				s = lca;
			}
		}
		// lca to t
		T right = ei;
		while (t != lca) {
			if (chain_id[t] != chain_id[lca]) {
				T tmp = query_in_chain(
					chain_id[t], 0, id_in_chain[t] - 1
				);
				right = merge(tmp, right);
				t = chain_top[chain_id[t]];
			}
			else {
				T tmp = query_in_chain(
					chain_id[t], id_in_chain[lca], id_in_chain[t] - 1
				);
				right = merge(tmp, right);
				t = lca;
			}
		}
		return merge(left, right);
	}

	/* edge update query */
	void update(int s, int t, T d) {
		if (depth[s] >= depth[t]) swap(s, t);
		if (chain_id[s] == chain_id[t])
			chain_update(chain_id[s], id_in_chain[s], d);
		else
			chain_update(chain_id[t], 0, d);
	}
};


int main() {
	int n; cin >> n;
	vector<vector<pair<int, Matrix>>> g(n);
	vector<pair<int, int>> edge(n - 1);
	for (int i = 0; i < n - 1; ++i) {
		int a, b; cin >> a >> b;
		edge[i] = { a, b };
		g[a].emplace_back(b, Matrix());
		g[b].emplace_back(a, Matrix());
	}
	HLD<Matrix> tree(g);
	tree.build();
	tree.show();
	int q; cin >> q;
	while (q--) {
		char c; cin >> c;
		if (c == 'x') {
			int i; ll x00, x01, x10, x11;
			cin >> i >> x00 >> x01 >> x10 >> x11;
			Matrix x;
			x.e00 = x00;
			x.e01 = x01;
			x.e10 = x10;
			x.e11 = x11;
			tree.update(edge[i].first, edge[i].second, x);
		}
		else {
			int i, j; cin >> i >> j;
			Matrix res = tree.edge_query(j, i, Matrix());
			cout << res.e00 << " ";
			cout << res.e01 << " ";
			cout << res.e10 << " ";
			cout << res.e11 << endl;
		}
	}
	return 0;
}
0