空夜 [换根DP]

空夜

Description

给定 \(n\) 个节点的树，每个点有点权 \(a_i\)，对于每个 \(i\)，求出 \(\sum_{j} \lfloor \frac{a_i}{2^{dis(i,j)}} \rfloor\)。

\(dis(i,j)\) 表示 \(i\) 到 \(j\) 的树上最短路径。

Solution

• 对于每个 \(i\) 都要求答案，等价于求以 \(i\) 为根的树的答案，可以想到 换根DP。

• 套路地，我们先 树形DP 求出以 \(1\) 为根的树的答案，然后再考虑换根。

• 设 \(f_u\) 表示以 \(u\) 为根的子树的答案，\(f_u=a_u+\sum (f_v\div 2)\)，\(v\) 是 \(u\) 的儿子。但这样是对的吗？由于有下取整，所以我们直接把 \(f_v\) 除以二显然是会把答案算大的。

• 我们从二进制上考虑，除以二下取整的本质是什么。显然是把该二进制右移一位，把第一位舍去了，那我们是不是可以记录这个子树的答案的第一位一共有多少个是 \(1\)。

• 但是只记录第一位转移不了，于是我们记录 \(g_{i,j}\) 表示以 \(i\) 为根的子树的答案的第 \(j\) 位有多少个 \(1\)。转移就是，\(g_{u,i}=\sum g_{v,i+1}\)。

• 那这时候 \(f\) 的转移就是 \(f_u=a_u+\sum((f_v-g_{v,1})\div 2)\)。

• 换根的转移比较简单，具体见代码。

Code

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
bool START_MEMORY;

const int N = 3e5 + 5, LogN = 40;

int n;
ll a[N], pow2[40];
ll dp[N], cnt[N][LogN]; // dp[u] 表示以 u 为根的子树的答案，cnt[i][j] 表示以 i 为根，总贡献的二进制第 j 位有多少个 1.

vector <int> G[N];

void dfs1(int u, int fa) {
    dp[u] = a[u];
    for (int v : G[u]) {
        if (v == fa) continue;
        dfs1(v, u);
        dp[u] += (dp[v] - cnt[v][0]) / 2ll;
        for (int i = 0; i <= 32; i++) cnt[u][i] += cnt[v][i + 1];
    }
}

void dfs2(int u, int fa) {
    for (int v : G[u]) {
        if (v == fa) continue;
        dp[v] += (dp[u] - (dp[v] - cnt[v][0]) / 2ll - (cnt[u][0] - cnt[v][1])) / 2;
        for (int i = 0; i <= 32; i++) cnt[v][i] += cnt[u][i + 1] - cnt[v][i + 2];
        dfs2(v, u);
    }
}

void Solve() {
    pow2[0] = 1;
    for (int i = 1; i <= 35; i++) pow2[i] = pow2[i - 1] * 2;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        for (int j = 0; j <= 32; j++) {
            if (a[i] & pow2[j]) cnt[i][j] = 1;
        }
    }
    int u, v;
    for (int i = 1; i < n; i++) {
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 0);
    for (int i = 1; i <= n; i++) cout << dp[i] << " ";
    cout << "\n";
}

bool END_MEMORY;
int main() {
    // freopen("sora.in", "r", stdin);
    // freopen("sora.out", "w", stdout);
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);

    Solve();

    cerr << "Time: " << 1000 * clock() / CLOCKS_PER_SEC << " ms\n";
    cerr << "Memory: " << fixed << setprecision(3) << double(&END_MEMORY - &START_MEMORY) / 1024 << " KB\n";
    return 0;
}