PROBLEM LINK
Author and Editorialist: Soumik Sarkar
Tester: Sarthak Manna
DIFFICULTY
Medium
PREREQUISITES
Dynamic programming on trees
PROBLEM
There exists a tree with n vertices where the cost of selecting the vertex i is a_i if no neighbour is selected and b_i otherwise, where b_i < a_i. Find maximum number of vertices that can be selected with total cost \leq k.
EXPLANATION
We will be using dynamic programming to solve this problem. Root the tree arbitrarily. Let us assume chosen root is vertex 1. Let f(i, j) be the minimum cost of selecting j vertices from the subtree rooted at i, then the answer is the maximum j for which f(1, j) \leq k holds.
We require 3 functions:
- f_0(i, j) is the cost of selecting j vertices from subtree of i such that i itself is not selected.
- f_a(i, j) is the cost of selecting j vertices from subtree of i such that i itself is selected but no child of i is selected.
- f_b(i, j) is the cost of selecting j vertices from subtree of i such that i itself and one or more of i's children are also selected.
Then we can say that f(i, j) = \min \{f_0(i, j), f_a(i, j), f_b(i, j)\}.
Let c_1, c_2, \dots c_k be the children of i. Then f_0(i, j) can be defined as
…which is the best way to distribute a total count of j over the k children.
However, calculating all ways to distribute j in this form is too inefficient. Here we can apply dynamic programming again over the children of i. You can read about this technique here (problem 3) and here.
To describe it in brief, f_0(i, \cdot), f_a(i, \cdot), and f_b(i, \cdot) are maintained as 1-D arrays. Iterating over the children of i, we consider all pairs of counts of previously processed vertices and the vertices in the subtree of the currently considered child. This is used to update new arrays which then replace the old arrays.
Pseudocode for the solution:
f0 = [0, INF]
fa = [INF, a[u]]
fb = [INF, INF]
sz[u] = 1
for each child v of u:
tmp0 = tmpa = tmpb = empty array of size sz[u] + sz[v] + 1
fill tmp0, tmpa, tmpb with INF
for i in [0..sz[u]]:
for j in [0..sz[v]]:
// Taking exactly i vertices from subtrees before v
// ...and exactly j vertices from subtree of v.
// For tmp0, any of f0[v], fa[v], fb[v] works
tmp0[i + j] = min(tmp0[i + j], f0[u][i] + min(f0[v][j], fa[v][j], fb[v][j]))
// For tmpa, child must not be selected, so only consider f0[v]
tmpa[i + j] = min(tmpa[i + j], fa[u][i] + f0[v][j])
// For tmpb,
// change fa to fb by selecting child v
// or keep fb and select/do not select v
tmpb[i + j] = min(tmpb[i + j],
fa[u][i] - a[u] + b[u] + min(fa[v][j] - a[v] + b[v], fb[v][j]),
fb[u][i] + min(f0[v][j], fa[v][j] - a[v] + b[v], fb[v][j]))
f0 = tmp0
fa = tmpa
tb = tmpb
sz[u] += sz[v]
Time complexity of this approach is \mathcal{O}(n^2).
Bonus problem
Instead of a_i and b_i, you are given all required c_{i,j} which is the cost of selecting vertex i such that exactly j neighbors of i are also selected. Can you find an \mathcal{O}(n^3) solution?
AUTHOR’S AND TESTER’S SOLUTIONS
Author’s solution can be found here
Tester’s solution can be found here.