In all solutions, the overall time complexity will be $Q$ times the time complexity to answer a single query.
First, let $I$ be the set of infected cows and let $T$ be the number of nights in a given query. We first analyze whether or not $I$ is even a valid configuration. Note that a cow $v$ can start out infected if and only if the set $B_T(v) = \{u \colon d(u, v) \leq T\}$ is contained entirely in $I$. Let $S$ be the set of cows that can start out infected. Then, $I$ is a valid configuration if and only if $I = \bigcup_{v\in S} B_T(v)$. Furthermore, when $I$ is a valid configuration, we want to find the smallest possible subset $S'\subseteq S$ such that $\bigcup_{v\in S'} B_T(v) = I$.
To compute the set $S$, we need to know the shortest distance from each cow to its closest non-infected cow. This can be done in $O(N)$ time by starting a multisource BFS from all non-infected cows.
Subtask 1 ($N = 10$):
For this case, we can just brute force through all subsets $S'$ of $S$. For each $v\in S'$, we can do a BFS to find $B_T(v)$, and check if the resulting sets cover $I$. This takes $O(2^N \cdot N^2)$ time per query.
Brandon's code:
N = int(input())
infected = [c == '1' for c in input()]
adj = [[] for _ in range(N)]
for _ in range(1, N):
u, v = [int(x) for x in input().split()]
adj[u-1].append(v-1)
adj[v-1].append(u-1)
dist = []
for u in range(N):
dist.append([N+1 for _ in range(N)])
bfs_queue = []
dist[u][u] = 0
bfs_queue.append(u)
bfs_queue_idx = 0
while len(bfs_queue) != bfs_queue_idx:
v = bfs_queue[bfs_queue_idx]
for w in adj[v]:
if dist[u][w] > dist[u][v] + 1:
dist[u][w] = dist[u][v] + 1
bfs_queue.append(w)
bfs_queue_idx += 1
def works(Sprime, T):
infected_by_Sprime = [False for _ in range(N)]
for u in Sprime:
for v in range(N):
if dist[u][v] <= T:
infected_by_Sprime[v] = True
for u in range(N):
if infected_by_Sprime[u] != infected[u]:
return False
return True
def query(T):
S = []
for u in range(N):
u_works = True
for v in range(N):
if dist[u][v] <= T and not infected[v]:
u_works = False
if u_works:
S.append(u)
if not works(S, T):
return -1
ans = len(S)
for i in range(2**len(S)):
Sprime = []
for j in range(len(S)):
if i // (2**j) % 2 == 1:
Sprime.append(S[j])
if works(Sprime, T):
ans = min(ans, len(Sprime))
return ans
Q = int(input())
for _ in range(Q):
T = int(input())
print(query(T))
Subtask 2 (all cows infected):
We can answer each query in $O(N)$ time using a greedy strategy. First, some definitions:
Our goal in this subtask is to choose the minimum number of cows so that all cows are covered.
Root the tree at any cow (the solution code roots the tree at the first cow). The following greedy procedure covers all the cows in $c$'s subtree at a distance at least $T$ away from $c$:
In addition, this procedure constructs an optimal solution because
Ben's code:
#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef pair<int, int> pi;
#define mp make_pair
#define f first
#define s second
#define sz(x) int((x).size())
#define rsz resize
#define pb push_back
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define F0R(i, a) FOR(i, 0, a)
#define rep(a) F0R(_, a)
vector<vi> adj;
string infected;
int N;
int n, ans;
// state:
// dist d to closest chosen: {1, d}
// dist d to farthest uncovered: {-1, d}
pi comb(pi a, pi b) {
if (a > b) swap(a, b);
if (b.f == -1) return max(a, b);
if (a.f == 1) return min(a, b);
assert(a.f == -1 && b.f == 1);
if (a.s + b.s > n) return a; // uncovered
return b; // covered
}
pi dfs(int x, int p) { // cover x's subtree
pi best{-1, 0};
for (int y : adj[x])
if (y != p) {
pi d = dfs(y, x);
++d.s;
best = comb(best, d);
}
if (best.f == -1 && (best.s == n || x == 0)) { // need to choose x
best = {1, 0};
++ans;
}
return best;
}
int query() {
ans = 0;
dfs(0, -1);
return ans;
}
int main() {
cin.tie(0)->sync_with_stdio(0);
cin >> N;
cin >> infected;
adj.rsz(N);
rep(N - 1) {
int a, b;
cin >> a >> b;
--a, --b;
adj.at(a).pb(b), adj.at(b).pb(a);
}
int Q;
cin >> Q;
rep(Q) {
cin >> n;
cout << query() << "\n";
}
}
Subtask 3 ($N\leq 400$) Solution 1 (DP):
Our approach for the remaining subtasks will be to use a tree-dp type approach. As in the subtask 2 solution, root the tree arbitrarily. Unfortunately, just making the DP state "min # of originally infected cows to cover subtree of $v$" does not provide enough info, because we can cover infected nodes in a subtree via starting cows outside the subtree.
Instead, for each $0 \leq t \leq T$, we store the following:
We can compute each of these in $O(\# \text{children})$ time given the children's DP states, so this takes a total of $O(TN)$ per query. Since $T\leq N$, this takes $O(N^2)$, which runs in time.
Subtask 3 Solution 2 (Greedy):
We will describe an algorithm where initially we start with no chosen cows, and we greedily choose cows in $S$ until every infected cow is covered.
Define $U$ to be the set of nodes that have not been covered yet. Root the tree arbitrarily and consider the highest depth infected cow $c \in U$, where depth is defined as distance from the root.
Let $S_c$ be all elements of $S$ that are distance within $T$ of $c$. Clearly, we must choose some element $s \in S_c$.
We claim that we can greedily choose the lowest depth node $s' \in S_c$. Consider some other node $s \in S_c$. For our greedy to be correct, it suffices to show that for all nodes $u \in U$, if $s$ covers $u$, then $s'$ also covers $u$.
This is easy to prove if $s$ and $s'$ are both ancestors of $c$. Now, suppose one of them is not an ancestor of $c$. Define $l = lca(c, s)$ and $l' = lca(c, s')$. Notice that $s$ covers all uncovered nodes $u \in U$ that are also in the subtree rooted at $l$ because $c$ is the highest depth uncovered node, and similarly $s'$ covers all uncovered nodes in the subtree rooted at $l'$. From this, one can prove that $s$ covers the same uncovered nodes as the unique node with the same depth as $s$ which is also an ancestor of $c$. Combining this with the proof for the case where $s$ and $s'$ are both ancestors of $c$, the greedy is proven.
So, the algorithm is to repeatedly find the highest depth node $c$ that has not been covered, and choose the lowest depth node in $S_c$ to cover it. This can be done naively in $O(N^2)$ per query (though on most test cases, it runs much faster than this worst-case bound would suggest).
Richard's code:
#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int mx = 100005;
int N;
int dist_to_nonsick[mx];
vector<int> adj[mx];
int depth[mx];
vector<int> sorted_by_depth;
string sick;
int k;
void genDepth(int node, int p = -1, int d = 0){
depth[node] = d;
for(auto u: adj[node]){
if(u == p) continue;
genDepth(u, node, d+1);
}
}
int dist_to_coverer[mx];
int searchForCoverer(int node, int p = -1, int d = 0){
pair<int, int> res = make_pair(MOD, -1);
if(dist_to_nonsick[node] > k){
res = min(res, make_pair(depth[node], node));
}
if(d+1 <= k){
for(auto u: adj[node]){
if(u == p) continue;
int other_res = searchForCoverer(u, node, d+1);
if(other_res != -1){
//keep track of minimum depth node within distance k
res = min(res, make_pair(depth[other_res], other_res));
}
}
}
return res.second;
}
void setCoverer(int node, int d = 0){
if(dist_to_coverer[node] <= d){
return;
}
dist_to_coverer[node] = d;
for(auto u: adj[node]){
setCoverer(u, d+1);
}
}
void solve(int _k){
k = _k;
for(int i = 1; i <= N; i++){
dist_to_coverer[i] = k+1; //minimum distance to a coverer
}
int ans = 0;
for(int i = int(sorted_by_depth.size())-1; i >= 0; i--){ //start with highest depth nodes
int to_cover = sorted_by_depth[i];
if(sick[to_cover] == '0') continue;
if(dist_to_coverer[to_cover] <= k) continue;
//this node must be covered
ans++;
int coverer = searchForCoverer(to_cover);
if(coverer == -1){ //can't be covered
ans = -1;
break;
}
//update dist_to_coverer
setCoverer(coverer);
}
cout << ans << "\n";
}
int main() {
cin.tie(0)->sync_with_stdio(0);
cin >> N;
cin >> sick; sick = "?" + sick;
for(int i = 1; i <= N-1; i++){
int a, b; cin >> a >> b;
adj[a].push_back(b);
adj[b].push_back(a);
}
// Compute minimum distance to a cow that is not sick with BFS
for(int i = 1; i <= N; i++){
dist_to_nonsick[i] = MOD;
}
queue<pair<int, int>> nonsick_q;
for(int i = 1; i <= N; i++){
if(sick[i] == '0'){
nonsick_q.push(make_pair(i, 0));
dist_to_nonsick[i] = 0;
}
}
while(nonsick_q.size()){
pair<int, int> a = nonsick_q.front(); nonsick_q.pop();
int node = a.first;
int d = a.second;
if(dist_to_nonsick[node] < d) continue;
for(auto u: adj[node]){
if(dist_to_nonsick[u] > d+1){
dist_to_nonsick[u] = d+1;
nonsick_q.push(make_pair(u, dist_to_nonsick[u]));
}
}
}
genDepth(1);
//sort nodes by depth
for(int i = 1; i <= N; i++) sorted_by_depth.push_back(i);
sort(begin(sorted_by_depth), end(sorted_by_depth), [&](int a, int b){
return depth[a] < depth[b];
});
int Q; cin >> Q;
for(int i = 1; i <= Q; i++){
int nights; cin >> nights;
solve(nights);
}
}
Full Solution 1:
We will answer each query in $O(N)$. This solution follows from generalizing the subtask 2 greedy solution or removing unnecessary states from the subtask 3 DP solution (it turns out that we need to keep at most two states for each subtree). For a node $v$, let $closest_S(v)$ denote the distance from $v$ to its closest node in $S$.
For any vertex $v$, let $I_v$ be the set of infected cows in $v$'s subtree, and let $S_v$ be the set of source cows in $v$'s subtree. In addition, let $I_v'$ be the set of infected cows in $v$'s subtree that are distance greater than $T-closest_S(parent_v)$ away from $parent_v$. When all cows are infected, $I_v'$ is precisely the set of all cows in $v$'s subtree distance at least $T$ away from $v$. Observe that every node in $I_v'$ must be covered by some node inside $v$'s subtree, and every node in $I_v\setminus I_v'$ would be covered by choosing the closest cow to $parent_v$. Intuitively, $I_v\setminus I_v'$ is the set of all cows in $v$'s subtree that might be covered by some node outside of $v$'s subtree (though they might instead be covered by some node inside $v$'s subtree).
Let $f(v)$ be the smallest number of nodes in $S_v$ we need to pick so that every node in $I_v'$ is covered. Clearly $f(v) \geq \sum_{u\in C_v} f(u)$. In fact the following is also true:
Lemma: $f(v) \leq 1 + \sum_{u\in C_v} f(u)$.
Proof: We will show that all of $I_v$ can be covered by at most $1 + \sum_{u\in C_v} f(u)$ source nodes (one of which need not be in $v$'s subtree). Choosing all of these source nodes (excluding any outside of $v$'s subtree) would cover $I_v'$.
Suppose that, in each subtree, we pick $f(u)$ sources in $S_u$ that cover $I_u'$, and also we are in the best possible DP state (states are defined in the solution above). Here, "every infected node in $u$'s subtree is covered, as is everything distance $T$ outside of $u$" is the best case, followed by "every infected node in $u$'s subtree is covered, as is everything distance $T-1$ outside of $u$", etc., and then "every infected node $u$'s subtree is covered, but nothing outside the subtree is" is better than "every infected node in $u$'s subtree is covered except $u$", which is better than "the deepest uncovered infected node is at depth 1", which is better than "the deepest uncovered infected node is at depth 2", etc.
Now, suppose $u_1, \ldots, u_k$ are the children such that $I_{u_1}, \ldots, I_{u_k}$ are not fully covered. Suppose their deepest uncovered children are at depth $d_1, \ldots, d_k$, respectively with respect to $v$. Without loss of generality suppose $d_1 = \max(d_1, \ldots, d_k)$. By assumption, there exists some source node $w$ such that $d(w, parent_{u_1})=d(w,v) \leq T - d_1$. Then, $d(w, v) \leq T-d_i$ for every $i$, so adding $w$ to the set of sources covers all of $I_v$. This completes the proof (note that $w$ covers $v$ if needed also). $\Box$
Since the answer to the query is just $f(\text{root})$, we just need to figure out when the difference $f(v) - \sum_{u\in C_v} f(u)$ is $0$ or $1$. To do this, when computing the DP, we additionally store what the "best case" is. We can do this in $O(\#\text{children})$ for each node with some preprocessing. Then, the difference is one precisely when there is an uncovered node in $I_v'$, in which case we can choose the closest node in $S$ to $v$. This newly chosen node is guaranteed to cover all uncovered nodes in $I_v$ (any nodes in $v$'s subtree not covered by this new node are part of $I_c'$ for some child $c$ of $v$, which should have been covered already). Note that due to our earlier observation about $I_v'$, this node will be in $v$'s subtree.
Ben's code:
#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef pair<int, int> pi;
#define mp make_pair
#define f first
#define s second
#define sz(x) int((x).size())
#define rsz resize
#define pb push_back
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define F0R(i, a) FOR(i, 0, a)
#define rep(a) F0R(_, a)
vector<vi> adj;
string infected;
vi closest_not_infected;
int N;
vi closest_ok;
int n, ans;
// state:
// dist d to closest chosen: {1, d}
// dist d to farthest uncovered: {-1, d}
// nothing: {0, 0}
pi comb(pi a, pi b) {
if (a > b) swap(a, b);
if (a.f == 0) return b;
if (b.f == 0) return a;
if (b.f == -1) return max(a, b);
if (a.f == 1) return min(a, b);
assert(a.f == -1 && b.f == 1);
if (a.s + b.s > n) return a; // uncovered
return b; // covered
}
pi dfs(int x, int p) { // cover x's subtree
pi best{infected.at(x) == '1' ? -1 : 0, 0};
for (int y : adj[x])
if (y != p) {
pi d = dfs(y, x);
if (d.f != 0) ++d.s;
best = comb(best, d);
}
if (best.f == -1) {
// check that choosing closest node in S to x would cover everything in
// x's subtree
assert(best.s + closest_ok.at(x) <= n);
if (p == -1 || best.s + 1 + closest_ok[p] > n) {
// there's something that must be covered by something in x's
// subtree
// choose closest node in S to x, which will be in x's
// subtree
++ans;
best = {1, closest_ok.at(x)};
}
}
return best;
}
vi multi_bfs(const vi &sources) {
vi ret(N, N + 1);
queue<int> q;
for (int i : sources) {
ret.at(i) = 0;
q.push(i);
}
while (sz(q)) {
int x = q.front();
q.pop();
for (int y : adj.at(x))
if (ret.at(x) + 1 < ret.at(y)) {
ret.at(y) = ret.at(x) + 1;
q.push(y);
}
}
return ret;
}
int query() {
vi origins; // S in the analysis
F0R(i, N) if (closest_not_infected.at(i) > n) origins.pb(i);
// distance from each node to closest element of S
closest_ok = multi_bfs(origins);
F0R(i, N)
if (infected.at(i) == '1' && closest_ok.at(i) > n)
return -1; // i can't be covered by anything in S
ans = 0;
dfs(0, -1);
return ans;
}
int main() {
cin.tie(0)->sync_with_stdio(0);
cin >> N;
cin >> infected;
adj.rsz(N);
rep(N - 1) {
int a, b;
cin >> a >> b;
--a, --b;
adj.at(a).pb(b), adj.at(b).pb(a);
}
vi sources;
F0R(i, N) if (infected.at(i) == '0') sources.pb(i);
closest_not_infected = multi_bfs(sources);
int Q;
cin >> Q;
rep(Q) {
cin >> n;
cout << query() << "\n";
}
}
Full Solution 2:
We can use centroid decomposition in the subtask 3 greedy solution to answer queries in $O(N\log{N})$.
Specifically, using centroid decomposition, we can do the following tasks.
We leave the details as an exercise to the reader.