#451690 (C++14) No.650 行列木クエリ

提出ソース

結果

問題	No.650 行列木クエリ
ユーザー	simkaren
提出日時	2020-03-30 20:26:15
言語	C++14 (gcc 12.3.0 + boost 1.83.0)
結果	TLE
実行時間	-
コード長	7,959 bytes
コンパイル時間	4,647 ms
コンパイル使用メモリ	267,852 KB
実行使用メモリ	47,988 KB
最終ジャッジ日時	2024-06-23 02:07:33
合計ジャッジ時間	8,552 ms
ジャッジサーバーID （参考情報）	judge4 / judge2

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 10,624 KB
testcase_01	TLE	-
testcase_02	--	-
testcase_03	--	-
testcase_04	--	-
testcase_05	--	-
testcase_06	--	-
testcase_07	--	-
testcase_08	--	-
testcase_09	--	-
testcase_10	--	-

権限があれば一括ダウンロードができます

ソースコード

raw source code

#pragma GCC optimize("Ofast", "unroll-loops")

#include <bits/stdc++.h>

using namespace std;

#define ll long long

// #define TEST

const ll mod = 1000000007LL;

struct Matrix {
	int n;
	vector<vector<ll>> element;

	Matrix() : n(2) {
		element.resize(n);
		for (int i = 0; i < n; ++i)
			element[i].resize(n, 0);
		for (int i = 0; i < n; ++i) element[i][i] = 1;
	}

	Matrix operator *(Matrix a) {
		assert(n == a.n);
		Matrix res;
		for (int i = 0; i < n; ++i) {
			for (int j = 0; j < n; ++j) {
				ll tmp = 0;
				for (int k = 0; k < n; ++k) {
					tmp += element[i][k] * a.element[k][j] % mod;
					tmp %= mod;
				}
				res.element[i][j] = tmp;
			}
		}
		return res;
	}

	Matrix pow(ll k) {
		if (k == 0) {
			Matrix res;
			for (int i = 0; i < n; ++i)
				res.element[i][i] = 1;
			return res;
		}
		Matrix tmp = pow(k / 2);
		if (k % 2 == 1)
			return tmp * tmp * (*this);
		return tmp * tmp;
	}
};

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) {
		int lca = LCA(s, t);
		if (chain_id[s] == chain_id[t]) {
			if (depth[s] <= depth[t])
				return query_in_chain(chain_id[s], id_in_chain[s], id_in_chain[t] - 1);
			else
				return query_in_chain_reverse(chain_id[s], id_in_chain[s] - 1, id_in_chain[t]);
		}
		// s to lca
		T left = s == lca ? T() : query_in_chain_reverse(
			chain_id[s], id_in_chain[s] - 1, 0
		);
		s = chain_top[chain_id[s]];
		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 = t == lca ? T() : query_in_chain(
			chain_id[t], 0, id_in_chain[t] - 1
		);
		t = chain_top[chain_id[t]];
		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.element[0][0] = x00;
			x.element[0][1] = x01;
			x.element[1][0] = x10;
			x.element[1][1] = x11;
			tree.update(edge[i].first, edge[i].second, x);
		}
		else {
			int i, j; cin >> i >> j;
			Matrix res = tree.edge_query(j, i);
			cout << res.element[0][0] << " ";
			cout << res.element[0][1] << " ";
			cout << res.element[1][0] << " ";
			cout << res.element[1][1] << endl;
		}
	}
	return 0;
}

yukicoder

結果

テストケース

ソースコード