TREESWAP1 - Editorial

editorial
, , , ,
1

PROBLEM LINK:

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

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

Combinatorics, DFS

PROBLEM:

You’re given a tree on N vertices.
For two binary strings A and B, define f(A, B) as follows:

  • You can choose an edge (u, v) of the tree and perform \text{swap}(A_u, A_v).
  • f(A, B) is the minimum number of operations needed to make A=B.
  • If it’s impossible to make A=B, then f(A, B) = 0.

You’re given the strings A and B.
However, some characters of A are replaced by ?.
Compute the sum of f(A', B) across all binary strings A' that can be obtained by filling in the blank spaces of A.

Here, N \le 5000.

EXPLANATION:

Let’s understand how to compute f(A, B) for fixed strings A and B.

We can only swap values around, so if A and B don’t have the same number of ones it’s clearly impossible to make them equal and hence f(A, B) = 0.
On the other hand, if A and B have the same number of ones, it’s always possible to make them equal via some sequence of swaps; so we focus on finding the minimum number.

Consider some edge (u, v) of the tree.
Deleting this edge splits the tree into two components. Let’s take one of these two components and call it P.
Now let:

  • P_A denote the number of vertices x such that x belongs to P and A_x = 1
  • P_B denote the number of vertices x such that x belongs to P and B_x = 1

Since we want A=B in the end, the counts of ones of A and B in P must be eventually equal.
The only operation that can change the count of ones in P is swapping along the (u, v) edge; since any other swap is either fully inside P or fully outside P.
Each swap can, at best, swap a 0 with a 1 (hence changing the count of ones inside P by 1.)

Thus, we definitely need at least |P_A - P_B| operations along the (u, v) edge on our way to make A and B equal.
(Observe that because A and B have the same count of ones, the value of |P_A - P_B| doesn’t depend on which of the components was chosen as P.)

There was nothing special about (u, v), so the above conclusion can be applied to every single edge of the tree.
Thus, a simple lower bound on f(A, B) is obtained by summing up |P_A - P_B| across all edges.

This lower bound is indeed attainable, and hence is the answer.

Proof

Let the cost of an edge be the number of swaps that need to be made along it.

We’ll show that as long as A \ne B and their counts of 1 are the same, there always exists an operation that reduces the cost of the edge it’s made along.

First, we delete all edges with cost 0 from the tree.
We are now left with several components. Since A \ne B, at least one of these components is not a singleton - choose any such non-singleton component.

Within the chosen component, each edge has a positive cost.
However, there’s also a fixed ‘direction’ we want to make swaps along, i.e. we always want to move a 1 from one specific side of the edge into the other side.

So, for each edge of the component, direct the edge from the side that is providing the 1’s to the side that needs the 1’s.

After directing all the edges, look at the component as a whole - it will be a directed acyclic graph.
Let s be any “source” vertex of this DAG (i.e. s is any vertex with zero indegree.)
From s, repeatedly follow edges till you reach some “sink” vertex t (so t has zero outdegree.)

Observe that we must have A_s = 1 (and B_s = 0), and similarly A_t = 0 and B_t = 1 as well.

Now consider the path we followed from s to t.
It started with a 1 and ended with a 0.
So, somewhere along the path, there must have been an edge u \to v such that A_u = 1 and A_v = 0.
Since the edge is directed u \to v, we want to swap ones from the side of u to the side of v anyway, along this edge.
So, simply performing the operation along this edge will reduce its cost!

This proves our claim of a valid operation always existing.
Simply repeating this over and over will make A = B eventually, with no wasted moves.

We can use this characterization of f(A, B) to solve the problem at hand.

Note that the answer for each edge is independent of all the others, allowing us to work with one edge at a time.
That is, we can fix one edge and attempt to sum up its contribution across all completions of A; and adding this up across all edges will give us the final answer.

So, suppose we fix the edge (u, v).
Let P denote the component containing u, after deleting edge (u, v), and Q denote the component containing v.

Note that the contribution of this edge is |P_A - P_B|, where P_B is fixed but P_A might not be.
Since N is small, we can thus try all possible values of P_A.

Now, suppose we fix the value of P_A (which must lie in [0, N]).
Then,

  • We need P_A ones in A present among the vertices of P.
    • Suppose there are x ones already present in P, and y empty spaces (i.e. y occurrences of '?')
    • Then, there are \binom{y}{P_A - x} ways to fill in the y empty spaces to end up with P_A ones here.
    • Note that if 0 \le P_A-x\le y doesn’t hold then this binomial coefficient is 0, which is fine.
  • However, we also need to maintain the condition of A and B having an equal number of ones in total.
    This requires C - P_A ones in the other tree Q, where C denotes the total number of ones present in B.
    • Once again, if there are x_2 ones already present in Q and y_2 empty spaces, the number of ways to fill them in to satisfy this is \binom{y_2}{C-P_A-x_2}.

So, upon fixing the value of P_A, the number of valid configurations can be found in \mathcal{O}(1) by multiplying a couple of binomial coefficients.
Add this count, multiplied by |P_A - P_B|, to the overall answer.

Since we do \mathcal{O}(N) work for each edge, this is \mathcal{O}(N^2) overall which is fast enough for this version of the problem.

To compute the number of ones and free positions present in each of the two smaller trees present after deleting an edge, you can either use subtree aggregates (since by rooting the tree at a fixed vertex, deleting an edge will always have one of the two parts be a rooted subtree), or just utilize the constraints being low and run a DFS/BFS in \mathcal{O}(N).

Binomial coefficients can be computed in \mathcal{O}(1) time with \mathcal{O}(N) precomputation via factorials and inverse factorials; alternately you can precompute and store all binomial coefficients with arguments \le N in \mathcal{O}(N^2) using Pascal’s identity.

TIME COMPLEXITY:

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

CODE:

Editorialist's code (C++) 
// #include <bits/allocator.h>
// #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);

    const int maxn = 5005;
    const int mod = 998244353;
    vector C(maxn, vector(maxn, 0));
    for (int n = 0; n < maxn; ++n) {
        C[n][0] = 1;
        for (int r = 1; r <= n; ++r) {
            C[n][r] = (C[n-1][r] + C[n-1][r-1]) % mod;
        }
    }
    auto get = [&] (int n, int r) {
        if (n < r or r < 0) return 0;
        return C[n][r];
    };
    
    int t; cin >> t;
    while (t--) {
        int n; cin >> n;
        vector adj(n, vector<int>());
        for (int i = 0; i < n-1; ++i) {
            int u, v; cin >> u >> v;
            adj[--u].push_back(--v);
            adj[v].push_back(u);
        }
        string a, b; cin >> a >> b;
        
        vector<array<int, 3>> subct(n);
        auto dfs = [&] (const auto &self, int u, int p) -> void {
            subct[u][0] = a[u] == '1';
            subct[u][1] = a[u] == '?';
            subct[u][2] = b[u] == '1';
            for (int v : adj[u]) if (v != p) {
                self(self, v, u);
                for (int i : {0, 1, 2}) subct[u][i] += subct[v][i];
            }
        };
        dfs(dfs, 0, 0);
        
        int ans = 0;
        int tot = ranges::count(b, '1');
        // for each edge
        // add difference between current ones and reqd ones to answer
        // under condition of count of ones in both strings being equal
        for (int u = 1; u < n; ++u) {
            for (int x = subct[u][0]; x <= n; ++x) {
                // x ones inside
                // tot-x outside
                // add abs(x-req[u]) to answer
                int ways = get(subct[u][1], x-subct[u][0]);
                ways = (1ll * ways * get(subct[0][1] - subct[u][1], tot-x-(subct[0][0] - subct[u][0]))) % mod;
                ans = (ans + 1ll * ways * abs(x-subct[u][2])) % mod;
            }
        }
        cout << ans << '\n';
    }
}
TREESWAP2 - Editorial