CLMTRO - Editorial

PROBLEM LINK

Practice
Contest

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

f_0(i, j) = \min_{x_1 + x_2 + \dots + x_k = j} \{f(c_1, x_1) + f(c_2, x_2) + \dots +f(c_k, x_k) \}

…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.

2 Likes

Can u plz explain how the complexity is O(n^2)

I mean for a node u,merging takes O(sz(u)^2),

So for a bamboo like tree(basically linear tree…),wont the complexity be O(N^3/6)?

1 Like

See Um_nik’s explanation here or tfg’s explanation here.
To but it briefly the number of steps is in the order of number of pairs of vertices, which is O(n^2).

1 Like

Ooh wow wow wow!!i didnt see that sz[u]=1…thanx a lot!!nice trick:)

There’s a typo in the pseudo code:

Instead of tmp0[i + j] = min(tmp0[i + j], f0[u][i] + min(f0[v][j] + fa[v][j], fb[v][j]))

it should be tmp0[i + j] = min(tmp0[i + j], f0[u][i] + min(f0[v][j], fa[v][j], fb[v][j]))

Thanks, fixed!