PROBLEM LINK:
Contest Division 1
Contest Division 2
Contest Division 3
Contest Division 4
Setter: Danny Boy
Tester: Harris Leung
Editorialist: Trung Dang
DIFFICULTY:
2526
PREREQUISITES:
PROBLEM:
You have a tree of N vertices rooted at vertex 1, where N is even. The i-th vertex has a binary value A_i written on it (i.e, each A_i is either 0 or 1).
A tree is said to be balanced if and only if the number of zeros written on its vertices equals the number of ones written on its vertices.
You can select some disjoint subtrees and flip the values on them (i.e, make all 0’s into 1’s and all 1’s into 0’s). Formally, you can pick any subset (possibly empty) of vertices \{v_1, v_2, \ldots, v_k\} such that v_i is not an ancestor of v_j for i \neq j, and flip the value of A_x for every vertex x that is contained in the subtree of some v_i.
Find a way of performing this operation that makes the entire tree balanced. It can be proved that the solution always exists.
EXPLANATION:
Let’s order the vertices by non-decreasing heights, and call this order O = [O_1, O_2, \dots, O_N], where obviously O_1 = 1. We have this following observation:
- Taking any suffix of O and inverting all nodes in this suffix is equivalent to inverting some disjoint subtrees.
To see how this observation work, suppose we are inverting the suffix O_i, O_{i + 1}, \dots, O_N. Let the height of O_i be d. Essentially, we are inverting all nodes with height at least d + 1, as well as some nodes with height d. This corresponds to inverting the subtree of those d-height nodes, as well as some subtrees of some d + 1-height nodes that is not covered by those d-height nodes.
With this observation, we create a new array P, where P_i is equal to the amount of 1-value nodes minus the amount of 0-value nodes (for sake of simplicity, we call this the balance value) among O_1, O_2, \dots, O_i. We can calculate P_i pretty easily (it’s sort of a prefix sum).
Note that initially, the balance value of the whole tree is just P_N, and we want this balance value to be 0. Inverting any suffix O_i, O_{i + 1}, \dots, O_N will change the balance value to P_{i - 1} - (P_N - P_{i - 1}) = 2 \cdot P_{i-1} - P_N, so we want to choose some i such that P_{i - 1} = \frac{P_N}{2}, then invert the suffix O_{[i..N]}. Luckily, using the intermediate value theorem, we can prove that there always exist such an i, and therefore the problem is solved.
TIME COMPLEXITY:
Time complexity is O(N) per test case if we use BFS to generate the order.
SOLUTION:
Setter's Solution
##include <bits/stdc++.h>
#define ll long long
#define int long long
#define fi first
#define se second
using namespace std;
void db() {cout << endl;}
template <typename T, typename ...U> void db(T a, U ...b) {cout << a << ' ', db(b...);}
#ifdef Cloud
#define file freopen("input.txt", "r", stdin), freopen("output.txt", "w", stdout)
#else
#define file ios::sync_with_stdio(false); cin.tie(0)
#endif
const int N = 3e5 + 1, mod = 998244353;
ll inf = 1ll << 60;
int d[N], par[N]{};
vector<int> g[N];
void dfs(int u, int p){
for (int i : g[u]) if (i != p){
d[i] = d[u] + 1;
par[i] = u;
dfs(i, u);
}
}
void solve(){
int n;
cin >> n;
for (int i = 1; i <= n; i++) g[i].clear();
int a[n + 1], sum = 0;
for (int i = 1; i <= n; i++){
cin >> a[i];
sum += a[i] == 0 ? -1 : 1;
}
for (int i = 0; i < n - 1; i++){
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
vector<pair<int, int>> v;
dfs(1, -1);
for (int i = 1; i <= n; i++) v.push_back({d[i], i});
sort(v.begin(), v.end());
int suf = 0, tar = n;
for (int i = n - 1; i >= 0; i--){
suf += a[v[i].se] == 0 ? -1 : 1;
if (sum - suf * 2 == 0) {
tar = i;
break;
}
}
bool flag[n + 1]{};
for (int i = tar; i < n; i++) flag[v[i].se] = 1;
vector<int> ans;
for (int i = 1; i <= n; i++) if (flag[i] and !flag[par[i]]) ans.push_back(i);
cout << ans.size() << '\n';
for (int &i : ans) cout << i << ' ';
cout << '\n';
}
signed main(){
file;
int t = 1;
cin >> t;
while (t--) solve();
}
Tester's Solution
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define fi first
#define se second
int n,m;
string s;
int a[300001];
vector<int>adj[300001];
int ptr=0;
int st[300001],en[300001];
int p[300001];
vector<int>ans;
void dfs(int id,int pp){
st[id]=++ptr;
p[ptr]=p[ptr-1]+a[id];
for(auto c:adj[id]){
if(c==pp) continue;
dfs(c,id);
}
en[id]=ptr;
}
void solve(int id,int p,int ql,int qr){
if(ql<=st[id] && en[id]<=qr){
ans.push_back(id);return;
}
if(st[id]>qr || en[id]<ql) return;
for(auto c:adj[id]){
if(c==p) continue;
solve(c,id,ql,qr);
}
}
void solve(){
cin >> n;ptr=0;ans.clear();
for(int i=1; i<=n ;i++){
cin >> a[i];
adj[i].clear();
}
for(int i=1; i<n ;i++){
int u,v;cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
dfs(1,0);
for(int i=0; i<n ;i++){
int res=p[i]+(n-i)-(p[n]-p[i]);
if(res==n/2){
solve(1,0,i+1,n);
cout << ans.size() << '\n';
for(auto c:ans) cout << c << ' ';
cout << '\n';
return;
}
}
}
int main(){
ios::sync_with_stdio(false);
int t;cin >> t;while(t--) solve();
}
Editorialist's Solution
#include <bits/stdc++.h>
using namespace std;
int main() {
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector<int> a(n);
vector<vector<int>> adj(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
for (int i = 1; i < n; i++) {
int u, v; cin >> u >> v; u--; v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
// return order, parent
auto BFS = [&](int src) {
vector<int> par(n, -1), ord;
// src is never chosen as a subtree, so we are safe to set par[src] = src
par[src] = src; ord.push_back(src);
for (int i = 0; i < n; i++) {
int u = ord[i];
for (int v : adj[u]) {
if (par[v] == -1) {
par[v] = u;
ord.push_back(v);
}
}
}
return make_pair(ord, par);
};
auto [ord, par] = BFS(0);
vector<int> prf(n);
for (int i = 0; i < n; i++) {
prf[i] = (i == 0 ? 0 : prf[i - 1]) + (a[ord[i]] == 0 ? 1 : -1);
}
for (int i = 0; i < n; i++) {
if (prf[i] == prf[n - 1] / 2) {
vector<bool> mrk(n);
for (int j = i + 1; j < n; j++) {
mrk[ord[j]] = true;
}
vector<int> ans;
for (int u = 0; u < n; u++) {
if (mrk[u] && !mrk[par[u]]) {
ans.push_back(u + 1);
}
}
cout << ans.size() << '\n';
for (int v : ans) {
cout << v << " ";
}
cout << '\n';
break;
}
}
}
}
