diff --git a/content/4_Gold/DP_Trees.problems.json b/content/4_Gold/DP_Trees.problems.json index 9c724e1c1e..09b4cfea50 100644 --- a/content/4_Gold/DP_Trees.problems.json +++ b/content/4_Gold/DP_Trees.problems.json @@ -123,8 +123,7 @@ "isStarred": false, "tags": ["Tree", "DP"], "solutionMetadata": { - "kind": "autogen-label-from-site", - "site": "AC" + "kind": "internal" } }, { diff --git a/solutions/gold/ac-SelectEdges.mdx b/solutions/gold/ac-SelectEdges.mdx new file mode 100644 index 0000000000..5e35a5c96c --- /dev/null +++ b/solutions/gold/ac-SelectEdges.mdx @@ -0,0 +1,92 @@ +--- +id: ac-SelectEdges +source: AC +title: Select Edges +author: Justin Ji +--- + +## Explanation + +To solve this problem, we write a tree DP. Consider that the only factor we +need to differentiate nodes on is whether or not we can connect a given node +to its parent. Thus, our DP state is the following: + +- $\texttt{dp}[u][0]$ is the best result in the subtree of $u$ if we can connect $u$ to its parent +- $\texttt{dp}[u][1]$ is the best result in the subtree of $u$ if we can't connect $u$ to its parent + +For a given node $u$, the "gain" we get from adding this node in is: + +$$ +(w + \texttt{dp}[u][0]) - \texttt{dp}[u][1] +$$ + +Thus, we want to use our $d[u]$ allowed edges on the nodes that have the most +benefit when adding an edge to them. Note that we handle the cases for +$\texttt{dp}[u][0])$ and $\texttt{dp}[u][1]$ pretty similarly. + +## Implementation + +**Time Complexity:** $\mathcal{O}(N\log{N})$ + + + + +```cpp +#include +using namespace std; + +using ll = long long; + +int main() { + int n; + cin >> n; + + vector d(n); + for (int &i : d) { cin >> i; } + + vector>> adj(n); + for (int i = 0; i < n - 1; i++) { + int u, v, w; + cin >> u >> v >> w; + u--, v--; + adj[u].push_back({v, w}); + adj[v].push_back({u, w}); + } + + // DP state is for the relation with its parent node + // {can connect, can't connect} + vector> dp(n); + const auto dfs = [&](int u, int p, auto self) -> void { + auto &[can_con, no_con] = dp[u]; + vector diffs; // diff between using edge and not using edge + + for (const auto &[v, w] : adj[u]) { + if (v == p) { continue; } + self(v, u, self); + can_con += dp[v][1]; + no_con += dp[v][1]; + if (d[v] > 0) { diffs.push_back(w + dp[v][0] - dp[v][1]); } + } + + int leftover = d[u]; + + sort(begin(diffs), end(diffs), greater()); + for (const ll &val : diffs) { + if (leftover == 0 || val < 0) { break; } + if (leftover > 1) { + can_con += val; + no_con += val; + } else { + no_con += val; + } + leftover--; + } + }; + dfs(0, -1, dfs); + + cout << dp[0][1] << endl; +} +``` + + +