PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Binary search, tree DP, interpolation
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.
Also count the number of arrangements that attain this minimum weakness.
EXPLANATION:
From the easy version, we know that finding the minimum possible weakness is doable in \mathcal{O}(N^2 \log N) time using dynamic programming.
Let this minimum value be M.
Our goal is now to count assignments that attain a weakness of M.
First, we require a slow solution.
Root the tree at vertex 1.
Define dp(u, x, y) as count of assignments to the subtree of u such that:
- all unsafe components have size \le M,
- x guards have been placed, and
- the unsafe component containing u has size y.
When merging a child into a parent, this information is fairly easy to update (in similar fashion to the easy-version DP that only computes the min weakness), by just considering all three states of the merged edge (do nothing, or add a guard in one of the two directions.)
The complexity of this is \mathcal{O}(N^4) - or more precisely, \mathcal{O}(K^2 N^2).
Just as in the easy version, it looks to be \mathcal{O}(K^2 N^3) but limiting iteration to the subtree sizes shaves off a linear factor.
This is, however, still too slow for our constraints - and needs to be optimized.
To further optimize our solution, clearly we need to reduce the information we store in the DP states; since it’s too costly to have cubic-order states in total.
This can be done by viewing the DP states via the lens of generating functions.
Specifically, observe that when merging a child into a parent, the number of guards will always add up (and maybe receive an extra +1, depending on what we do with the connecting edge.)
This corresponds to degrees adding when multiplying polynomials; and maybe multiplying by z as well (where z is the polynomial variable.)
Thus, it makes sense to try and represent things as polynomials with the number of guards as the degree.
To not have to deal with any casework during the transitions, or rather to make things work nicely with polynomial algebra, it’s helpful to try and separate out a couple of states in the DP.
In particular, note that when merging states, the casework depends primarily on whether or not the root of the subtree is safe or not.
So, keep separate DP tables for whether u is safe or not.
Note that when u is unsafe we still do need to know the size of the unsafe component containing it; so there are actually several polynomials in each node.
After doing this, the transitions between states can all be represented somewhat nicely using purely polynomial algebra.
Now, at first glance this doesn’t seem particularly useful - it’s just a different way of looking at things, but nothing has actually been optimized yet.
However, this is where we can use the fact that we’re working with polynomial algebra.
Observe that the value we want to actually find, is the sum of coefficients of z^K in all the polynomials of the root.
The sum of polynomials is also just a polynomial itself so we’re looking for the coefficient of z^K in a single polynomial.
In particular, we don’t actually care about any of the intermediate polynomials at all - they just happen to be there during computation but their actual coefficients don’t quite matter to us.
We can thus use polynomial interpolation to find our target value.
That is, observe that since all the DP transitions involve simple polynomial algebra, this same polynomial algebra allows us to pick any value and evaluate the polynomial there too - without actually needing to store the coefficients!
So, we can pick N distinct values in [0, 998244353) and evaluate the polynomial at those values - and then interpolate those values to get the coefficients of the polynomial at the root.
This is fine because K \lt N.
Each polynomial evaluation via the DP takes \mathcal{O}(N^2) time, so this can be done in \mathcal{O}(N^3) overall which is fast enough.
TIME COMPLEXITY:
\mathcal{O}(N^3) 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());
/**
* Integers modulo p, where p is a prime
* Source: Aeren (modified from tourist?)
* Modmul for 64-bit mod from kactl:ModMulLL
* Works with p < 7.2e18 with x87 80-bit long double, and p < 2^52 ~ 4.5e12 with 64-bit
*/
template<typename T>
struct Z_p{
using Type = typename decay<decltype(T::value)>::type;
static vector<Type> MOD_INV;
constexpr Z_p(): value(){ }
template<typename U> Z_p(const U &x){ value = normalize(x); }
template<typename U> static Type normalize(const U &x){
Type v;
if(-mod() <= x && x < mod()) v = static_cast<Type>(x);
else v = static_cast<Type>(x % mod());
if(v < 0) v += mod();
return v;
}
const Type& operator()() const{ return value; }
template<typename U> explicit operator U() const{ return static_cast<U>(value); }
constexpr static Type mod(){ return T::value; }
Z_p &operator+=(const Z_p &otr){ if((value += otr.value) >= mod()) value -= mod(); return *this; }
Z_p &operator-=(const Z_p &otr){ if((value -= otr.value) < 0) value += mod(); return *this; }
template<typename U> Z_p &operator+=(const U &otr){ return *this += Z_p(otr); }
template<typename U> Z_p &operator-=(const U &otr){ return *this -= Z_p(otr); }
Z_p &operator++(){ return *this += 1; }
Z_p &operator--(){ return *this -= 1; }
Z_p operator++(int){ Z_p result(*this); *this += 1; return result; }
Z_p operator--(int){ Z_p result(*this); *this -= 1; return result; }
Z_p operator-() const{ return Z_p(-value); }
template<typename U = T>
typename enable_if<is_same<typename Z_p<U>::Type, int>::value, Z_p>::type &operator*=(const Z_p& rhs){
#ifdef _WIN32
uint64_t x = static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value);
uint32_t xh = static_cast<uint32_t>(x >> 32), xl = static_cast<uint32_t>(x), d, m;
asm(
"divl %4; \n\t"
: "=a" (d), "=d" (m)
: "d" (xh), "a" (xl), "r" (mod())
);
value = m;
#else
value = normalize(static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value));
#endif
return *this;
}
template<typename U = T>
typename enable_if<is_same<typename Z_p<U>::Type, int64_t>::value, Z_p>::type &operator*=(const Z_p &rhs){
uint64_t ret = static_cast<uint64_t>(value) * static_cast<uint64_t>(rhs.value) - static_cast<uint64_t>(mod()) * static_cast<uint64_t>(1.L / static_cast<uint64_t>(mod()) * static_cast<uint64_t>(value) * static_cast<uint64_t>(rhs.value));
value = normalize(static_cast<int64_t>(ret + static_cast<uint64_t>(mod()) * (ret < 0) - static_cast<uint64_t>(mod()) * (ret >= static_cast<uint64_t>(mod()))));
return *this;
}
template<typename U = T>
typename enable_if<!is_integral<typename Z_p<U>::Type>::value, Z_p>::type &operator*=(const Z_p &rhs){
value = normalize(value * rhs.value);
return *this;
}
template<typename U>
Z_p &operator^=(U e){
if(e < 0) *this = 1 / *this, e = -e;
Z_p res = 1;
for(; e; *this *= *this, e >>= 1) if(e & 1) res *= *this;
return *this = res;
}
template<typename U>
Z_p operator^(U e) const{
return Z_p(*this) ^= e;
}
Z_p &operator/=(const Z_p &otr){
Type a = otr.value, m = mod(), u = 0, v = 1;
if(a < (int)MOD_INV.size()) return *this *= MOD_INV[a];
while(a){
Type t = m / a;
m -= t * a; swap(a, m);
u -= t * v; swap(u, v);
}
assert(m == 1);
return *this *= u;
}
template<typename U> friend const Z_p<U> &abs(const Z_p<U> &v){ return v; }
Type value;
};
template<typename T> bool operator==(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value == rhs.value; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator==(const Z_p<T>& lhs, U rhs){ return lhs == Z_p<T>(rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator==(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) == rhs; }
template<typename T> bool operator!=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return !(lhs == rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator!=(const Z_p<T> &lhs, U rhs){ return !(lhs == rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator!=(U lhs, const Z_p<T> &rhs){ return !(lhs == rhs); }
template<typename T> bool operator<(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value < rhs.value; }
template<typename T> bool operator>(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value > rhs.value; }
template<typename T> bool operator<=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value <= rhs.value; }
template<typename T> bool operator>=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value >= rhs.value; }
template<typename T> Z_p<T> operator+(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator+(const Z_p<T> &lhs, U rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator+(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T> Z_p<T> operator-(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator-(const Z_p<T>& lhs, U rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator-(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T> Z_p<T> operator*(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator*(const Z_p<T>& lhs, U rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator*(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T> Z_p<T> operator/(const Z_p<T> &lhs, const Z_p<T> &rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator/(const Z_p<T>& lhs, U rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator/(U lhs, const Z_p<T> &rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T> istream &operator>>(istream &in, Z_p<T> &number){
typename common_type<typename Z_p<T>::Type, int64_t>::type x;
in >> x;
number.value = Z_p<T>::normalize(x);
return in;
}
template<typename T> ostream &operator<<(ostream &out, const Z_p<T> &number){ return out << number(); }
/*
using ModType = int;
struct VarMod{ static ModType value; };
ModType VarMod::value;
ModType &mod = VarMod::value;
using Zp = Z_p<VarMod>;
*/
// constexpr int mod = 1e9 + 7; // 1000000007
constexpr int mod = (119 << 23) + 1; // 998244353
// constexpr int mod = 1e9 + 9; // 1000000009
using Zp = Z_p<integral_constant<decay<decltype(mod)>::type, mod>>;
template<typename T> vector<typename Z_p<T>::Type> Z_p<T>::MOD_INV;
template<typename T = integral_constant<decay<decltype(mod)>::type, mod>>
void precalc_inverse(int SZ){
auto &inv = Z_p<T>::MOD_INV;
if(inv.empty()) inv.assign(2, 1);
for(; inv.size() <= SZ; ) inv.push_back((mod - 1LL * mod / (int)inv.size() * inv[mod % (int)inv.size()]) % mod);
}
template<typename T>
vector<T> precalc_power(T base, int SZ){
vector<T> res(SZ + 1, 1);
for(auto i = 1; i <= SZ; ++ i) res[i] = res[i - 1] * base;
return res;
}
template<typename T>
vector<T> precalc_factorial(int SZ){
vector<T> res(SZ + 1, 1); res[0] = 1;
for(auto i = 1; i <= SZ; ++ i) res[i] = res[i - 1] * i;
return res;
}
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) {
// already safe -> remains safe
if (i == 0) {
if (j <= m) ndp[0] = min(ndp[0], dp[u][i] + dp[v][j]);
}
else {
// not safe -> 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 orient edge
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;
};
auto getmin = [&] () {
int lo = 0, hi = n;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (check(mid)) hi = mid;
else lo = mid + 1;
}
return lo;
};
int M = getmin();
vector dp(n, vector<Zp>());
auto calc = [&] (const auto &self, int u, int p, int x) -> void {
dp[u] = {0, 1};
for (auto [v, w] : adj[u]) {
if (v == p) continue;
self(self, v, u, x);
int tot = dp[u].size() + subsz[v];
vector ndp(tot, Zp(0));
for (int i = 0; i < ssize(dp[u]); ++i) for (int j = 0; j <= subsz[v]; ++j) {
if (i == 0) {
// no guard -> j must be valid
if (j <= M) ndp[0] += dp[u][0] * dp[v][j];
if (w == 0) { // place guard here
// guard towards u -> j must be valid
if (j <= M) ndp[0] += x * dp[u][0] * dp[v][j];
// guard towards v -> no constraint
ndp[0] += x * dp[u][0] * dp[v][j];
}
}
else {
// no guard -> j must be valid
if (j <= M) ndp[i+j] += dp[u][i] * dp[v][j];
if (w == 0) { // place guard here
// towards u -> j must be valid
if (j <= M) ndp[0] += x * dp[u][i] * dp[v][j];
// towards v -> no constraint
ndp[i] += x * dp[u][i] * dp[v][j];
}
}
}
dp[u] = move(ndp);
}
};
vector vals(n, Zp(0));
for (int i = 0; i < n; ++i) {
calc(calc, 0, -1, i);
for (int j = 0; j <= M; ++j) vals[i] += dp[0][j];
}
// interpolate to find coef of x^k
// https://github.com/kth-competitive-programming/kactl/blob/main/content/numerical/PolyInterpolate.h
auto interpolate = [&] (auto x, auto y, int N) {
vector<Zp> res(N), temp(N);
for (int j = 0; j < N-1; ++j) for (int i = j+1; i < N; ++i) {
y[i] = (y[i] - y[j]) / (x[i] - x[j]);
}
Zp last = 0; temp[0] = 1;
for (int j = 0; j < N; ++j) for (int i = 0; i < N; ++i) {
res[i] += y[j] * temp[i];
swap(last, temp[i]);
temp[i] -= last * x[j];
}
return res;
};
vector pts(n, Zp(0));
for (int i = 0; i < n; ++i) pts[i] = i;
auto poly = interpolate(pts, vals, n);
cout << M << ' ' << poly[k] << '\n';
}
}