题意
略。
题解
完全为了水博客
处理出每个点bitset
的后缀或(从所有儿子结点继承来)和每个点的bitset
的前缀或(从父亲的方向继承来)。
bitset
以路径长度为关键字。这样以一个点为起点的路径长度种数就是这个点上前后缀bitset
或的count()
。
这样可以做到(mathcal O(frac {n ^ 2}{omega}))。
但是要注意的是,一个点的前缀或必须和其兄弟节点一起在其父亲节点处处理(还是处理一个前缀和数组),否则复杂度就不对了。
#pragma GCC optimize(2)
#include <bits/stdc++.h>
using namespace std;
const int N = 5e4 + 5;
int n, a[N];
vector <int> g[N];
bitset <N> tmp, pre[N], suf[N], s[N];
void dfs (int x, int p) {
suf[x][0] = a[x];
for (auto y : g[x]) {
if (y != p) {
dfs(y, x);
suf[x] |= suf[y] << 1;
}
}
}
void sfd (int x, int p) {
vector <int> ch(0);
tmp = pre[x], tmp[0] = a[x];
for (auto y : g[x]) {
if (y != p) {
ch.push_back(y);
}
}
s[ch.size()].reset();
for (int i = (int)ch.size() - 1; i >= 0; --i) {
s[i] = s[i + 1] | (suf[ch[i]] << 1);
}
for (int i = 0; i < (int)ch.size(); ++i) {
pre[ch[i]] = (tmp | s[i + 1]) << 1;
tmp |= suf[ch[i]] << 1;
}
for (auto y : ch) {
sfd(y, x);
}
}
int main () {
scanf("%d", &n);
for (int i = 1; i <= n; ++i) {
scanf("%d", &a[i]);
}
for (int i = 1, x, y; i < n; ++i) {
scanf("%d%d", &x, &y);
g[x].push_back(y), g[y].push_back(x);
}
dfs(1, 0), sfd(1, 0);
for (int i = 1; i <= n; ++i) {
printf("%d
", (int)(pre[i] | suf[i]).count());
}
return 0;
}