二乗の木 DP の例題として、すごくいい感じ!!!
問題概要
頂点のツリーが与えられて、各頂点 には 人がいる。いま、人を移動させて、各頂点の人数の最大値と最小値との差が 1 になるようにしたい。
1 人が辺 1 個分を移動するのに必要なコストが 1 である。これを達成するのに必要な総コストの最小値を求めよ。
制約
考えたこと
総合人数を all 人とすると
- all/N 人のノード数
- all/N + 1 人のノード数
がそれぞれどうなるべきかがわかる。K = all%N とすると、N 頂点のうち適切に K 個を選んで、そこに all/N + 1 人にするのに必要な最小コストを求める問題ということになる。
さて、こういうのはいかにも DP っぽい
- dp[v][k] := v を根とする熱つき木において、余り頂点数を k 個の状態にするのに必要な最小コスト (ただし、人材過多や人材不足は v の親ノードで調節する)
という風にすればよさそう。一見 かかるが、二乗の木DP の構造になっていて でできる。
#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; }