2023 rcpc

2023/rcpc/d2/editorial/m-dragons-694.cpp

@@ -0,0 +1,393 @@
+#include <bits/stdc++.h>
+
+using namespace std;
+
+#ifdef LOCAL
+ifstream fin("dragons.in");
+ofstream fout("dragons.out");
+#else
+#define fin cin
+#define fout cout
+#endif
+
+void YN(bool cond) {
+	if(cond) {
+		fout << "YES\n";
+	} else {
+		fout << "NO\n";
+	}
+}
+
+const int kV = 1000;
+const int kQ = 1e4;
+const int kN = 1e4;
+const int kLog = 14;
+const int kS = 1000;
+const int kB = /*2 * kN / sqrt(kQ)*/ 120;
+
+int weight[kN];
+vector<int> plusminus;
+int tin[kN], tout[kN];
+vector<int> adj[kN];
+int par[kN][kLog + 1], rmq[kN][kLog + 1];
+
+void dfs(int u = 0, int v = 0) {
+	plusminus.emplace_back(u);
+	tin[u] = plusminus.size() - 1;
+
+	par[u][0] = v;
+	rmq[u][0] = weight[u];
+	for(int i = 1; i <= kLog; i++) {
+		par[u][i] = par[par[u][i - 1]][i - 1];
+		rmq[u][i] = max(rmq[u][i - 1], rmq[par[u][i - 1]][i - 1]);
+	}
+
+	for(const auto &it: adj[u]) if(it != v) {
+		dfs(it, u);
+	}
+
+	plusminus.emplace_back(u);
+	tout[u] = plusminus.size() - 1;
+}
+
+bool upper(int u, int v) {
+	return tin[u] <= tin[v] && tin[v] <= tout[u];
+}
+
+int lcaQuery(int u, int v) {
+	if(upper(u, v)) {
+		return u;
+	}
+	if(upper(v, u)) {
+		return v;
+	}
+
+	for(int i = kLog; i >= 0; i--) {
+		if(par[u][i] && !upper(par[u][i], v)) {
+			u = par[u][i];
+		}
+	}
+
+	return par[u][0];
+}
+
+void maxSelf(int &x, int y) {
+	if(y > x) {
+		x = y;
+	}
+}
+
+int maxQuery(int u, int v) {
+	if(upper(u, v)) {
+		swap(u, v);
+	}
+
+	int ret = 0;
+	for(int i = kLog; i >= 0; i--) {
+		if(par[u][i] && !upper(par[u][i], v)) {
+			maxSelf(ret, rmq[u][i]);
+			u = par[u][i];
+		}
+	}
+	maxSelf(ret, rmq[u][0]);
+
+	return ret;
+}
+
+int maxChain(int u, int v) {
+	int lca = lcaQuery(u, v), res = weight[lca];
+	maxSelf(res, maxQuery(u, lca));
+	maxSelf(res, maxQuery(v, lca));
+	return res;
+}
+
+struct Query {
+	int left, right, mx, hp, idx, lca, bucket;
+
+	Query() {}
+	Query(int left, int right, int mx, int hp, int idx, int lca): left(left), right(right), mx(mx), hp(hp), idx(idx), lca(lca), bucket(left / kB) {}
+
+	bool operator < (const Query &oth) const {
+		if(mx == oth.mx) {
+			if(bucket == oth.bucket) {
+				if(bucket & 1) {
+					return right > oth.right;
+				}
+				return right < oth.right;
+			}
+			return bucket < oth.bucket;
+		}
+		return mx < oth.mx;
+	}
+};
+
+vector<Query> Qs;
+vector<bool> ans;
+
+struct subsetSum {
+	struct bol {
+		uint32_t v1;
+		uint64_t v2;
+		// __uint128_t v3;
+
+		bol() {}
+		bol(uint32_t v1, uint64_t v2/*, __uint128_t v3*/): v1(v1), v2(v2)/*, v3(v3)*/ {}
+		
+		bol& operator += (const bol &addValue) {
+			v1 += addValue.v1;
+			v2 += addValue.v2;
+			// v3 += addValue.v3;
+			return *this;
+		}
+		
+		bol& operator += (const uint64_t &addValue) {
+			v1 += addValue;
+			v2 += addValue;
+			// v3 += addValue;
+			return *this;
+		}
+		
+		bol& operator -= (const uint64_t &subValue) {
+			v1 -= subValue;
+			v2 -= subValue;
+			// v3 -= subValue;
+			return *this;
+		}
+		
+		bol& operator -= (const bol &subValue) {
+			v1 -= subValue.v1;
+			v2 -= subValue.v2;
+			// v3 -= subValue.v3;
+			return *this;
+		}
+		
+		bool operator == (bool cond) const {
+			return (v1 || v2/* || v3*/) == cond;
+		}
+		
+		bool operator ! () const {
+			return !(v1 || v2/* || v3*/);
+		}
+	};
+	int n;
+	vector<bol> dp;
+
+	subsetSum() {}
+	subsetSum(int n): n(n), dp(n) {
+		// printf("subsetSum dp(%d);\n", n);
+		init();
+	}
+
+	void init() {
+		dp[0] = bol(1, 1);
+		for(int i = 1; i < n; i++) {
+			dp[i] = bol(0, 0);
+		}
+	}
+
+	void insert(int weight) {
+		// printf("dp.insert(%d);\n", weight);
+		for(int i = n - 1; i >= weight; i--) { 
+			dp[i] += dp[i - weight];
+		}
+	}
+
+	void erase(int weight) {
+		// printf("dp.erase(%d);\n", weight);
+		for(int i = weight; i < n; i++) {
+			dp[i] -= dp[i - weight];
+		}
+	}
+
+	bool query(int weight) {
+		// printf("cout << dp.query(%d) << '\\n';\n", weight);
+		return dp[weight] == 1;
+	}
+
+	bool operator [] (int weight) {
+		return query(weight);
+	}
+};
+
+subsetSum dp;
+vector<bool> marked;
+int frq[kV + 1], to_update[kV + 1];
+
+void add(int u) {
+	if(marked[u] == 0) {
+		frq[weight[u]]++;
+		to_update[weight[u]]++;
+		marked[u] = 1;
+	} else {
+		frq[weight[u]]--;
+		to_update[weight[u]]--;
+		marked[u] = 0;
+	}
+}
+
+void updateAll(int mx) {
+	for(int i = 1; i < mx; i++) {
+		if(to_update[i] < 0) {
+			to_update[i] = -to_update[i];
+			while(to_update[i]) {
+				to_update[i]--;
+				dp.erase(i);
+			}
+		} else {
+			while(to_update[i]) {
+				to_update[i]--;
+				dp.insert(i);
+			}
+		}
+	}
+}
+
+void updateOne(int i) {
+	int cnt = 1;
+	if(to_update[i] < 0) {
+		to_update[i] = -to_update[i];
+		while(to_update[i]) {
+			to_update[i]--;
+			dp.erase(i);
+		}
+	} else {
+		while(to_update[i]) {
+			to_update[i]--;
+			dp.insert(i);
+		}
+	}
+}
+
+int main() {
+	fin.tie(nullptr)->sync_with_stdio(false);
+	int n;
+	assert(fin >> n);
+	assert(2 <= n && n <= kN);
+
+	for(int i = 0; i < n; i++) {
+		assert(fin >> weight[i]);
+		assert(1 <= weight[i] && weight[i] <= kV);
+	}
+
+	for(int i = 1; i < n; i++) {
+		int u, v;
+		assert(fin >> u >> v);
+		assert(1 <= u && u <= n);
+		assert(1 <= v && v <= n);
+		u--; v--;
+
+		adj[u].emplace_back(v);
+		adj[v].emplace_back(u);
+	}
+
+	dfs();
+
+	int q;
+	assert(fin >> q);
+	assert(1 <= q && q <= kQ);
+	ans = vector<bool>(q);
+
+	for(int i = 0; i < q; i++) {
+		int u, v, hp;
+		assert(fin >> u >> v >> hp);
+		assert(1 <= u && u <= n);
+		assert(1 <= v && v <= n);
+		assert(0 <= hp && hp <= kV);
+		u--; v--;
+
+		if(tin[u] > tin[v]) {
+			swap(u, v);
+		}
+
+		int lca = lcaQuery(u, v), mx = maxChain(u, v);
+
+		if(lca == u) {
+			Qs.emplace_back(tin[u], tin[v], mx, hp, i, -1);
+		} else {
+			Qs.emplace_back(tout[u], tin[v], mx, hp, i, lca);
+		}
+	}
+
+	sort(Qs.begin(), Qs.end());
+	dp = subsetSum(kS + 1);
+
+	int lst = -1, l, r;
+
+	for(int i = 0; i < q; i++) {
+		int left = Qs[i].left;
+		int right = Qs[i].right;
+		int mx = Qs[i].mx;
+		int hp = Qs[i].hp;
+		int lca = Qs[i].lca;
+
+		if(mx != lst) {
+			memset(to_update, 0, sizeof(to_update));
+			memset(frq, 0, sizeof(frq));
+			marked = vector<bool>(n);
+			dp.init();
+
+			l = left;
+			r = right;
+
+			for(int j = l; j <= r; j++) {
+				add(plusminus[j]);
+			}
+		} else {
+			while(l < left) {
+				add(plusminus[l]); // rem
+				l++;
+			}
+
+			while(l > left) {
+				l--;
+				add(plusminus[l]);
+			}
+
+			while(r < right) {
+				r++;
+				add(plusminus[r]);
+			}
+
+			while(r > right) {
+				add(plusminus[r]); // rem
+				r--;
+			}
+		}
+
+		if(lca != -1) {
+			add(lca);
+		}
+		updateAll(mx);
+
+		int frq_mx = frq[mx];
+
+		bool crt_ans = 0;
+		for(int cnt = 0; crt_ans == 0 && cnt <= frq_mx && cnt * mx <= hp; cnt++) {
+			int crt_hp = hp - cnt * mx, crt_frq = frq_mx - cnt;
+			if(crt_frq & 1) {
+				if(crt_hp >= mx && dp[crt_hp - mx]) {
+					crt_ans = 1;
+				}
+			} else if(dp[crt_hp]){
+				crt_ans = 1;
+			}
+		}
+
+		ans[Qs[i].idx] = crt_ans;
+
+		if(lca != -1) {
+			add(lca); // rem
+			if(weight[lca] != mx) {
+				updateOne(weight[lca]);
+			}
+		}
+		to_update[mx] = 0;
+
+		lst = mx;
+	}
+
+	for(int i = 0; i < q; i++) {
+		YN(ans[i]);
+	}
+	return 0;
+}