#include <bits/stdc++.h>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;
template <class T> using Tree = tree<T, null_type, less<T>, rb_tree_tag,tree_order_statistics_node_update>;
template <class T> struct LazySeg {
std::vector<T> sum, lazy;
int sz;
void init(int sz_) {
sz = 1;
while (sz < sz_) sz *= 2;
sum.assign(2 * sz, 0);
lazy.assign(2 * sz, 0);
}
void push(int ind, int L, int R) {
sum[ind] += (R - L + 1) * lazy[ind];
if (L != R) {
lazy[2 * ind] += lazy[ind];
lazy[2 * ind + 1] += lazy[ind];
}
lazy[ind] = 0;
}
void pull(int ind) {
sum[ind] = sum[2 * ind] + sum[2 * ind + 1];
}
void build() {
for (int i = sz - 1; i >= 1; i--) {
pull(i);
}
}
void upd(int lo, int hi, T inc, int ind = 1, int L = 0, int R = -1) {
if (R == -1) R += sz;
push(ind, L, R);
if (hi < L || R < lo) return ;
if (lo <= L && R <= hi) {
lazy[ind] = inc;
push(ind, L, R);
return;
}
int M = (L + R) / 2;
upd(lo, hi, inc, 2 * ind, L, M);
upd(lo, hi, inc, 2 * ind + 1, M + 1, R);
pull(ind);
}
T qsum(int lo, int hi, int ind = 1, int L = 0, int R = -1) {
if (R == -1) R += sz;
push(ind, L, R);
if (lo > R || L > hi) return 0;
if (lo <= L && R <= hi) return sum[ind];
int M = (L + R) / 2;
return qsum(lo, hi, 2 * ind, L, M) + qsum(lo, hi, 2 * ind + 1, M + 1, R);
}
};
int main() {
// ios_base::sync_with_stdio(0);
int n, k;
cin >> n >> k;
vector<vector<int>> adj(n);
for (int i = 0; i < n - 1; i++) {
int u, v; cin >> u >> v;
u--, v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
vector<int> s(n);
for (int i = 0; i < n; i++)
cin >> s[i], s[i]--;
if (n == 1) {
cout << 1 << '\n';
return 0;
}
vector<int> L(n), R(n), parent(n);
int ti = 0;
function<int(int, int)> dfs_precomp = [&](int src, int par) -> int {
int mx = ti++;
parent[src] = par;
L[src] = mx;
for (int nxt : adj[src]) {
if (nxt == par)
continue;
mx = max(mx, dfs_precomp(nxt, src));
}
R[src] = mx;
return mx;
};
dfs_precomp(0, -1);
LazySeg<int> seg;
seg.init(n);
vector<int> oc(n), tot(n);
vector<Tree<int>> num(n);
for (int i = 0; i < n; i++)
num[s[i]].insert(L[i]);
for (int x : s)
tot[x]++;
function<void(int, int)> dfs_edges = [&](int src, int par) {
for (int nxt : adj[src]) {
if (nxt == par)
continue;
dfs_edges(nxt, src);
}
int val = num[s[src]].order_of_key(R[src] + 1);
val -= num[s[src]].order_of_key(L[src]);
if (val == 1) {
// update the subtree
seg.upd(L[src], R[src], 1);
seg.upd(L[src], L[src], -1);
}
if (val == tot[s[src]]) {
// update above
seg.upd(L[0], R[0], 1);
seg.upd(L[src], R[src], -1);
seg.upd(L[src], L[src], 1);
}
};
dfs_edges(0, -1);
vector<int> bad(n);
int amt = 0;
for (int i = 1; i < n; i++)
if (seg.qsum(L[i], L[i]) == k)
bad[i] = 1, amt++;
int leaves = 0;
function<int(int, int)> dfs_sub = [&](int src, int par) -> int {
int sub = bad[src];
for (int nxt : adj[src]) {
if (nxt == par)
continue;
sub += dfs_sub(nxt, src);
}
if (sub == 1 && bad[src])
leaves++;
else if (sub == amt && bad[src])
leaves++;
return sub;
};
dfs_sub(0, -1);
int ans = (leaves + 1) / 2;
cout << ans << '\n';
return 0;
}
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
187 ms |
262148 KB |
Execution killed with signal 9 (could be triggered by violating memory limits) |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
187 ms |
262148 KB |
Execution killed with signal 9 (could be triggered by violating memory limits) |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
248 ms |
262148 KB |
Execution killed with signal 9 (could be triggered by violating memory limits) |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
187 ms |
262148 KB |
Execution killed with signal 9 (could be triggered by violating memory limits) |
| 2 |
Halted |
0 ms |
0 KB |
- |
