けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 3048 Averaging (AUPC 2018 day2-J)

二乗の木 DP の例題として、すごくいい感じ!!!

問題へのリンク

問題概要

 N 頂点のツリーが与えられて、各頂点  v には  X_v 人がいる。いま、人を移動させて、各頂点の人数の最大値と最小値との差が 1 になるようにしたい。

1 人が辺 1 個分を移動するのに必要なコストが 1 である。これを達成するのに必要な総コストの最小値を求めよ。

制約

  •  2 \le N \le 5000

考えたこと

総合人数を all 人とすると

  • all/N 人のノード数
  • all/N + 1 人のノード数

がそれぞれどうなるべきかがわかる。K = all%N とすると、N 頂点のうち適切に K 個を選んで、そこに all/N + 1 人にするのに必要な最小コストを求める問題ということになる。

さて、こういうのはいかにも DP っぽい

  • dp[v][k] := v を根とする熱つき木において、余り頂点数を k 個の状態にするのに必要な最小コスト (ただし、人材過多や人材不足は v の親ノードで調節する)

という風にすればよさそう。一見  O(N^{3}) かかるが、二乗の木DP の構造になっていて  O(N^{2}) でできる。

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int N;
vector<long long> X;
vector<vector<int>> G;
long long all;
long long ave;
int amari;

int size[5100];
long long sum[5100];
long long dp[5100][5100];
long long sdp[5100][5100];

void rec(int v, int p = -1) {
    vector<int> chs;
    size[v] = 1;
    sum[v] = X[v];
    for (auto ch : G[v]) {
        if (ch == p) continue;
        rec(ch, v);
        chs.push_back(ch);
        size[v] += size[ch];
        sum[v] += sum[ch];
    }

    for (int i = 0; i <= chs.size(); ++i)
        for (int j = 0; j <= size[v]; ++j)
            sdp[i][j] = INF;
    sdp[0][0] = 0;
    int cur = 0;
    for (int i = 0; i < chs.size(); ++i) {
        int ch = chs[i];
        for (int j = 0; j <= cur; ++j) {
            for (int k = 0; k <= size[ch]; ++k) {
                chmin(sdp[i+1][j+k], sdp[i][j] + dp[ch][k]);
            }
        }
        cur += size[ch];
    }

    for (int k = 0; k <= size[v]; ++k) {
        chmin(dp[v][k], sdp[chs.size()][k] + abs(ave * size[v] + k - sum[v]));
        chmin(dp[v][k+1], sdp[chs.size()][k] + abs(ave * size[v] + k+1 - sum[v]));
    }
}

int main() {
    cin >> N;
    X.resize(N);
    all = 0;
    for (int i = 0; i < N; ++i) cin >> X[i], all += X[i];
    ave = all / N;
    amari = all % N;
    G.assign(N, vector<int>());
    for (int i = 0; i < N-1; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    for (int v = 0; v < N; ++v) for (int i = 0; i <= N; ++i) dp[v][i] = INF;
    rec(0);

    cout << dp[0][amari] << endl;
}