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:

Tester's code (C++) 
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif

/*



*/

const int MOD = 1e9 + 7;
const int N = 2e3 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

vector<pll> adj[N];
vector<bool> blocked(N);

ll dp[N][N][2]; // (node, size_of_unsafe, par_edge_used) = min required cuts
vector<ll> subsiz(N);
ll M;

void dfs1(ll u, ll p){
    subsiz[u] = 1;
    ll par_edge_id = -1;
    for(auto [v,id] : adj[u]){
        if(v == p){
            par_edge_id = id;
            conts;
        }
        dfs1(v,u);
        subsiz[u] += subsiz[v];
    }

    ll siz = subsiz[u];
    ll dp1[siz+5][2], dp2[siz+5][2];
    memset(dp1,0x3f,sizeof dp1);
    memset(dp2,0x3f,sizeof dp2);

    ll su = 1;
    dp1[su][0] = 0;

    for(auto [v,id] : adj[u]){
        if(v == p) conts;

        ll sv = subsiz[v];

        rep(s1,su+1){
           rep(s2,sv+1){
               rep(f1,2){
                   rep(f2,2){
                       ll val = dp1[s1][f1] + dp[v][s2][f2];
                       ll s3 = s1+s2;
                       ll f3 = f1;
                       if(!f2){
                           // par edge not used --> can use up only if actually free
                           if(!blocked[id]){
                               f3 = 1;
                           }
                       }

                       amin(dp2[s3][f3],val);
                   }
               }
           } 
        }

        su += sv;

        rep(i,su+1){
            rep(j,2){
                dp1[i][j] = dp2[i][j];
                dp2[i][j] = inf2;
            }
        }
    }

    // update actual big dp
    rep(i,su+1){
        rep(j,2){
            dp[u][i][j] = inf2;
        } 
    }
    
    rep(i,su+1){
        rep(j,2){
            // no extra cut
            if(i <= M){
                amin(dp[u][i][0],dp1[i][j]);
            }

            // make extra cut (at curr node)
            if(j == 0){
                // j = 0: no fr ee edge avail
                // can cut only if par edge is present and is not blocked
                if(par_edge_id != -1 and !blocked[par_edge_id]){
                    amin(dp[u][0][1],dp1[i][j]+1);
                }
            }
            else{
                // j = 1: free edge avail
                // no need to use par edge
                amin(dp[u][0][0],dp1[i][j]+1);
            }
        }
    }
}

void solve(int test_case){
    ll n,k; cin >> n >> k;

    rep1(i,n){
        adj[i].clear();
    }
    
    rep1(i,n-1){
        ll u,v,t; cin >> u >> v >> t;
        adj[u].pb({v,i});
        adj[v].pb({u,i});
        blocked[i] = (bool)t;
    }

    auto ok = [&](ll mid) -> bool{
        M = mid;
        dfs1(1,-1);
        
        rep(i,n+1){
            rep(j,2){
                if(dp[1][i][j] <= k){
                    return true;
                }
            }
        }

        return false;
    };
    
    ll lo = 1, hi = n-1;
    ll ans = n;
    while(lo <= hi){
        ll mid = (lo+hi)>>1;
        if(ok(mid)){
            ans = mid;
            hi = mid-1;
        }
        else{
            lo = mid+1;
        }
    }

    cout << ans << endl;
}

int main()
{
    fastio;

    int t = 1;
    cin >> t;

    rep1(i, t) {
        solve(i);
    }

    cerr << "RUN SUCCESSFUL" << endl;

    return 0;
}