Submission #64389607

---

Source Code Expand

#include <random>
#include <string>
#include <vector>
#include <cassert>
#include <iostream>
#include <algorithm>
using namespace std;

string to_string(const vector<int>& arr) {
	string res = "[";
	for (int i = 0; i < int(arr.size()); i++) {
		if (i != 0) {
			res += ", ";
		}
		res += to_string(arr[i]);
	}
	res += "]";
	return res;
}

struct parade {
	int a, b, l, r;
};

struct query {
	int l, r;
};

class mergesort_tree {
private:
	int size;
	vector<vector<pair<int, uint64_t> > > tsub;
	vector<vector<uint64_t> > txor;
	vector<int> T_;
	vector<uint64_t> H_;
public:
	mergesort_tree(int N, const vector<int>& T, const vector<uint64_t>& H) {
		size = 1 << (N >= 2 ? 32 - __builtin_clz(N - 1) : 0);
		tsub.resize(2 * size);
		for (int i = 0; i < N; i++) {
			tsub[size + i] = {make_pair(T[i], H[i])};
		}
		for (int i = size - 1; i >= 1; i--) {
			tsub[i].resize(tsub[i * 2].size() + tsub[i * 2 + 1].size());
			merge(tsub[i * 2].begin(), tsub[i * 2].end(), tsub[i * 2 + 1].begin(), tsub[i * 2 + 1].end(), tsub[i].begin());
		}
		txor.resize(2 * size);
		for (int i = 1; i < 2 * size; i++) {
			txor[i].resize(tsub[i].size() + 1);
			for (int j = 0; j < int(tsub[i].size()); j++) {
				txor[i][j + 1] = txor[i][j] ^ tsub[i][j].second;
			}
		}
		// for (int i = size - 1; i >= 1; i--) {
		// 	if (tsub[i].size() >= 2) {
		// 		for (int j = 1; j < tsub[i].size(); j++) {
		// 			tsub[i][j].second ^= tsub[i][j - 1].second;
		// 		}
		// 	}
		// }
		T_ = T;
		H_ = H;
	}
	pair<int, uint64_t> query(int l, int r, int d, int u) {
		/*
		int cnt = 0; uint64_t hx = 0;
		for (int i = l; i < r; i++) {
			if (d <= T_[i] && T_[i] < u) {
				cnt++;
				hx ^= H_[i];
			}
		}
		return make_pair(cnt, hx);
		*/
		pair<int, uint64_t> ans = {0, uint64_t(0)};
		auto check = [&](int x) -> void {
			int pd = lower_bound(tsub[x].begin(), tsub[x].end(), make_pair(d, uint64_t(0))) - tsub[x].begin();
			int pu = lower_bound(tsub[x].begin(), tsub[x].end(), make_pair(u, uint64_t(0))) - tsub[x].begin();
			// for (int i = pd; i < pu; i++) {
			// 	ans.first++;
			// 	ans.second ^= tsub[x][i].second;
			// }
			// uint64_t hd = (pd >= 1 ? tsub[x][pd - 1].second : 0);
			// uint64_t hu = (pu >= 1 ? tsub[x][pu - 1].second : 0);
			ans.first += pu - pd;
			ans.second ^= txor[x][pd] ^ txor[x][pu];
		};
		l += size;
		r += size;
		while (l != r) {
			if (l & 1) {
				check(l++);
			}
			if (r & 1) {
				check(--r);
			}
			l /= 2;
			r /= 2;
		}
		return ans;
	};
};

vector<bool> solve_path(int N, int M, int Q, const vector<int>& T, const vector<int>& C, const vector<parade>& P, const vector<query>& qs) {
	// step #1. hashing
	mt19937_64 mt(202503311258);
	vector<uint64_t> hbase(N);
	for (int i = 0; i < N; i++) {
		hbase[i] = mt();
	}
	vector<uint64_t> hc(N);
	for (int i = 0; i < N; i++) {
		hc[i] = hbase[C[i]];
	}

	// step #2. calculate for each parade
	mergesort_tree seg(N, T, hc);
	vector<uint64_t> hparade(M);
	for (int i = 0; i < M; i++) {
		if (P[i].a == P[i].b) {
			hparade[i] = ((P[i].r - P[i].l) % 2 == 0 ? uint64_t(0) : hc[P[i].a]);
		} else {
			pair<int, uint64_t> res1 = seg.query(0, P[i].a + 1, P[i].l, P[i].r);
			pair<int, uint64_t> res2 = seg.query(P[i].b, N, P[i].l, P[i].r);
			pair<int, uint64_t> res3 = seg.query(P[i].a + 1, P[i].b, P[i].l, P[i].r);
			// cerr << "res1: " << res1.first << ' ' << res1.second << endl;
			// cerr << "res2: " << res2.first << ' ' << res2.second << endl;
			// cerr << "res3: " << res3.first << ' ' << res3.second << endl;
			uint64_t h = res3.second;
			if (res1.first & 1) {
				h ^= hc[P[i].a];
			}
			if (res2.first & 1) {
				h ^= hc[P[i].b];
			}
			hparade[i] = h;
		}
	}

	// step #3. prefix sum
	vector<uint64_t> hprefix(M + 1);
	for (int i = 0; i < M; i++) {
		hprefix[i + 1] = hprefix[i] ^ hparade[i];
	}
	vector<uint64_t> hsort = hbase;
	sort(hsort.begin(), hsort.end());

	// step #4. solve
	vector<bool> ans(Q);
	for (int i = 0; i < Q; i++) {
		int l = qs[i].l, r = qs[i].r;
		uint64_t h = hprefix[l] ^ hprefix[r];
		ans[i] = (h == 0 || binary_search(hsort.begin(), hsort.end(), h));
	}

	return ans;
}

vector<bool> solve(int N, int M, int Q, const vector<vector<int> >& G, const vector<int>& C, const vector<parade>& P, const vector<query>& qs) {
	assert(N <= 2000 && M <= 2000);

	// step #1. build tree
	vector<int> par(N, -1), d(N);
	auto build_tree = [&](auto& self, int x, int pre) -> void {
		for (int i : G[x]) {
			if (i != pre) {
				par[i] = x;
				d[i] = d[x] + 1;
				self(self, i, x);
			}
		}
	};
	build_tree(build_tree, 0, -1);

	// step #2. get each "nearest"
	vector<vector<int> > tbl(M, vector<int>(N));
	for (int i = 0; i < M; i++) {
		vector<bool> path(N, false);
		int a = P[i].a, b = P[i].b;
		path[a] = true;
		path[b] = true;
		while (a != b) {
			if (d[a] < d[b]) {
				swap(a, b);
			}
			a = par[a];
			path[a] = true;
		}
		auto dfs = [&](auto& self, int x, int pre, int c) -> void {
			tbl[i][x] = c;
			for (int i : G[x]) {
				if (i != pre) {
					self(self, i, x, c);
				}
			}
		};
		for (int j = 0; j < N; j++) {
			if (path[j]) {
				tbl[i][j] = j;
				for (int k : G[j]) {
					if (!path[k]) {
						dfs(dfs, k, j, j);
					}
				}
			}
		}
	}

	// step #3. hashing
	mt19937_64 mt(202503311258);
	vector<uint64_t> hbase(N);
	for (int i = 0; i < N; i++) {
		hbase[i] = mt();
	}
	vector<uint64_t> hparade(M);
	for (int i = 0; i < M; i++) {
		for (int j = P[i].l; j < P[i].r; j++) {
			hparade[i] ^= hbase[C[tbl[i][j]]];
		}
	}
	vector<uint64_t> hprefix(M + 1);
	for (int i = 0; i < M; i++) {
		hprefix[i + 1] = hprefix[i] ^ hparade[i];
	}
	vector<uint64_t> hsort = hbase;
	sort(hsort.begin(), hsort.end());

	// step #4. solve
	vector<bool> ans(Q);
	for (int i = 0; i < Q; i++) {
		int l = qs[i].l, r = qs[i].r;
		uint64_t h = hprefix[l] ^ hprefix[r];
		ans[i] = (h == 0 || binary_search(hsort.begin(), hsort.end(), h));
	}

	return ans;
}

int main() {
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	int N;
	cin >> N;
	vector<int> A(N - 1), B(N - 1);
	for (int i = 0; i < N - 1; i++) {
		cin >> A[i] >> B[i];
		A[i]--; B[i]--;
	}
	bool is_path = true;
	for (int i = 0; i < N - 2; i++) {
		is_path = is_path && (B[i] == A[i + 1]);
	}
	vector<int> C(N);
	for (int i = 0; i < N; i++) {
		cin >> C[i];
		C[i]--;
	}
	int M;
	cin >> M;
	vector<parade> P(M);
	for (int i = 0; i < M; i++) {
		cin >> P[i].a >> P[i].b >> P[i].l >> P[i].r;
		P[i].a--; P[i].b--; P[i].l--;
	}
	int Q;
	cin >> Q;
	vector<query> qs(Q);
	for (int i = 0; i < Q; i++) {
		cin >> qs[i].l >> qs[i].r;
		qs[i].l--;
	}
	if (is_path) {
		vector<int> T = A;
		T.push_back(B[N - 2]);
		vector<int> TINV(N);
		for (int i = 0; i < N; i++) {
			TINV[T[i]] = i;
		}
		for (parade& p : P) {
			p.a = TINV[p.a];
			p.b = TINV[p.b];
			if (p.a > p.b) {
				swap(p.a, p.b);
			}
		}
		vector<int> TC(N);
		for (int i = 0; i < N; i++) {
			TC[i] = C[T[i]];
		}
		vector<bool> ans = solve_path(N, M, Q, T, TC, P, qs);
		for (int i = 0; i < Q; i++) {
			cout << (ans[i] ? "Yes" : "No") << '\n';
		}
	} else {
		vector<vector<int> > G(N);
		for (int i = 0; i < N - 1; i++) {
			G[A[i]].push_back(B[i]);
			G[B[i]].push_back(A[i]);
		}
		vector<bool> ans = solve(N, M, Q, G, C, P, qs);
		for (int i = 0; i < Q; i++) {
			cout << (ans[i] ? "Yes" : "No") << '\n';
		}
	}
	return 0;
}

Submission Info

Judge Result