2023 rcpc

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

@@ -0,0 +1,122 @@
+#include <bits/stdc++.h>
+
+using namespace std;
+
+#ifdef LOCAL
+ifstream fin("brute.in");
+ofstream fout("brute.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 kS = 1000;
+
+vector<int> depth, parent, weight;
+vector<vector<int>> adj;
+
+void dfs(int u = 0, int v = -1) {
+	for(const auto &it: adj[u]) if(it != v) {
+		depth[it] = depth[u] + 1;
+		parent[it] = u;
+		dfs(it, u);
+	}
+}
+
+bool check_greedy(const vector<int> &h, int hp) {
+	bitset<kS + 1> dp;
+	dp[0] = 1;
+	int mx = 0, frq = 1;
+	for(const auto &it: h) {
+		if(it > mx) {
+			mx = it;
+			frq = 1;
+		} else if(it == mx) {
+			frq++;
+		}
+	}
+	for(const auto &it: h) {
+		if(it != mx && it <= hp) {
+			dp |= dp << it;
+		}
+	}
+	for(int i = 0; i <= frq; i++) {
+		if(i * mx > hp) {
+			return 0;
+		}
+		int crt_hp = hp - i * mx, crt_frq = frq - i;
+		if(crt_frq & 1) {
+			if(crt_hp >= mx && dp[crt_hp - mx]) {
+				return 1;
+			}
+		} else {
+			if(dp[crt_hp]) {
+				return 1;
+			}
+		}
+	}
+	return 0;
+}
+
+int main() {
+	int n;
+	fin >> n;
+	adj = vector<vector<int>>(n);
+	weight = vector<int>(n);
+
+	for(auto &it: weight) {
+		fin >> it;
+	}
+
+	for(int i = 1; i < n; i++) {
+		int u, v;
+		fin >> u >> v;
+		u--; v--;
+
+		adj[u].emplace_back(v);
+		adj[v].emplace_back(u);
+	}
+
+	depth = vector<int>(n);
+	parent = vector<int>(n);
+	dfs();
+
+	int q;
+	fin >> q;
+	for(int i = 0; i < q; i++) {
+		int u, v, hp;
+		fin >> u >> v >> hp;
+		u--; v--;
+
+		vector<int> h;
+		if(depth[u] < depth[v]) {
+			swap(u, v);
+		}
+		while(depth[u] > depth[v]) {
+			h.emplace_back(weight[u]);
+			u = parent[u];
+		}
+		while(u != v) {
+			h.emplace_back(weight[u]);
+			h.emplace_back(weight[v]);
+			u = parent[u];
+			v = parent[v];
+		}
+		h.emplace_back(weight[u]); 
+		bool ans = check_greedy(h, hp);
+		YN(ans);
+	}
+
+	return 0;
+}