TREEGUARD - Editorial

editorial
, , , , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy-Medium

PREREQUISITES:

Binary search, tree DP

PROBLEM:

You’re given a tree with N vertices, and K.
Some edges of the tree are blocked.
You can place guards on any K non-blocked edges, and orient each of them towards one endpoint of the edge.

The weakness of an assignment is the maximum size of a connected subset of vertices such that no guard is oriented toward any of the chosen vertices.

Find the minimum possible weakness of the tree.

EXPLANATION:

When trying to minimize the maximum of several quantities, it’s often useful to try and explicitly fix the upper bound on allowed values, and then see if all quantities can be brought within this upper bound.
If we’re able to do this, the answer is the smallest upper bound that works, which can be found with an additional log factor using binary search.

So, let’s fix the value M, and see if we can make all unsafe components have size \le M using (at most) K guards.
(Using fewer than K guards is fine, since adding more guards cannot increase the weakness.)

Because the constraints allow for quadratic time, this is not too hard using dynamic programming on subtrees.

Let’s root the tree at vertex 1, and build up information from smaller subtrees to larger ones.
When merging children subtrees into a parent, we need to know the following piece of information:

  • For each child, what’s the size of the unsafe component containing it?

With this in mind, define dp(u, x) to be the minimum number of guards needed in the subtree of u such that:

  • All unsafe components within this subtree have size \le M, and
  • The unsafe component containing u has size x.

In particular, x=0 here means vertex u is safe.

Transitions are fairly simple: let our state be (u, x) and a child have state (v, y). Then, merging the child into u,

  • If x = 0 so u is already safe, merging anything into it will make it remain safe.
    The resulting state is (u, 0) with a cost of dp(u, x) + dp(v, y).
  • if x \gt 0, then merging (v, y) into it while not doing anything with the edge connecting them results in a state of (u, x+y), still with a cost of dp(u, x) + dp(v, y).
    • Note that this option is only available if x+y \le M.
  • Finally, if the edge between u and v is not blocked, we have the option to orient it in either direction.
    • If we orient towards v, the resulting state is (u, x) with a cost of dp(u, x) + dp(v, y) + 1.
    • If we orient towards u, the resulting state is (u, 0) since u becomes safe; with a cost again of dp(u, x) + dp(v, y) + 1.

Computing this DP looks like \mathcal{O}(N^3) because there are \mathcal{O}(N^2) states and \mathcal{O}(N) transitions from each.
However, it is actually \mathcal{O}(N^2) if implemented properly (i.e. by limiting loops to subtree sizes) - see point 7 of this blog.

The minimum number of guards needed to achieve this is then given by the minimum value at some valid state of the root (vertex 1, for us.)
If this is \le K then M is a valid upper bound, otherwise it’s too small.

This, combined with binary search, gives a solution in \mathcal{O}(N^2 \log N) time which is fast enough for the constraints.

TIME COMPLEXITY:

\mathcal{O}(N^2 \log N) per testcase.

CODE:

Editorialist's code (C++) 
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int t; cin >> t;
    while (t--) {
        int n, k; cin >> n >> k;
        
        vector adj(n, vector<array<int, 2>>());
        for (int i = 0; i < n-1; ++i) {
            int u, v, w; cin >> u >> v >> w;
            adj[--u].push_back({--v, w});
            adj[v].push_back({u, w});
        }

        vector subsz(n, 0);
        auto predfs = [&] (const auto &self, int u, int p) -> void {
            subsz[u] = 1;
            for (auto [v, w] : adj[u]) {
                if (v == p) continue;
                self(self, v, u);
                subsz[u] += subsz[v];
            }
        };
        predfs(predfs, 0, -1);

        auto check = [&] (int m) {
            vector dp(n, vector<int>());

            auto dfs = [&] (const auto &self, int u, int p) -> void {
                dp[u] = {k+1, 0};

                for (auto [v, w] : adj[u]) {
                    if (v == p) continue;
                    self(self, v, u);

                    int tot = dp[u].size() + subsz[v];
                    vector ndp(tot, k+1);

                    for (int i = 0; i < ssize(dp[u]); ++i) for (int j = 0; j <= subsz[v]; ++j) {
                        // u already safe -> remains safe
                        if (i == 0) {
                            if (j <= m) ndp[0] = min(ndp[0], dp[u][i] + dp[v][j]);
                        }
                        else {
                            // u not safe, v not safe, no guard -> unsafe size increases by j
                            if (j <= m) ndp[i+j] = min(ndp[i+j], dp[u][i] + dp[v][j]);
    
                        }
                        // not blocked -> try to place guard
                        if (!w) {
                            ndp[i] = min(ndp[i], dp[u][i] + dp[v][j] + 1); // towards v
                            if (j <= m) ndp[0] = min(ndp[0], dp[u][i] + dp[v][j] + 1); // towards u
                        }
                    }

                    dp[u] = move(ndp);
                }
            };
            dfs(dfs, 0, -1);

            return *min_element(begin(dp[0]), begin(dp[0])+m+1) <= k;
        };

        int lo = 0, hi = n;
        while (lo < hi) {
            int mid = (lo + hi) / 2;
            if (check(mid)) hi = mid;
            else lo = mid + 1;
        }

        cout << lo << '\n';
    }
}