换根DP的基本写法

换根DP的基本写法

树DP中，许多解法的思路都是子树向上合并。举例说明，无根树中，需要遍历k个点，求以各节点为根时最小代价。定义dp[u][x]为遍历完子树后，是否向上回溯的代价。写法可如下：

int sz[N], go[N];
int dp[N][2];
set< pair<int, int> > s[N];

void link (int u, int v, int w) {
	sz[u] += sz[v];
	if (sz[v]) {
		s[u].insert({ (dp[v][1] + 2 * w) - (dp[v][0] + w), v });
		dp[u][1] += dp[v][1] + 2 * w;
	}
	dp[u][0] = (sz[u] != go[u]) ? dp[u][1] - (--s[u].end())->first : 0;
}

void cut (int u, int v, int w) {
	sz[u] -= sz[v];
	if (sz[v]) {
		s[u].erase({ (dp[v][1] + 2 * w) - (dp[v][0] + w), v });
		dp[u][1] -= dp[v][1] + 2 * w;
	}
	dp[u][0] = (sz[u] != go[u]) ? dp[u][1] - (--s[u].end())->first : 0;
}

void dfs (int u, int from) {
	for (auto [v, w] : u) if (u != from) {
		dfs(v, u);
		link(u, v, w);
	}
}

void solve (int u, int from) {
	ans[u] = dp[u][0];
	for (auto [v, w] : u) if (u != from) {
		cut(u, v, w); link(v, u, w);
		solve(v, u);
		cut(v, u, w); link(u, v, w);
	}
}

思路非常简单，复杂度O(NlogN)，显然不是只维护最值达到的最低复杂度。我们发现每次仅仅取set的最值，每次仅仅断一条边，也就是可以只维护最大值和次大值，复杂度是O(N)，一种方便写法如下：

void link (int u, int v, int w) {
	...
	if (s[u].size() > 2) 
		s[u].erase(--s[u].end());
	...
}

不管看着爽不爽，这样写确实很短很方便。