空夜 [換根DP]

chenwenmo發表於2024-11-09

空夜

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;
}

相關文章