PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Combinatorics, Mo’s algorithm
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 2\cdot 10^5.
EXPLANATION:
To recap the solution from the easy version:
- f(A, B) = 0 if A and B have different counts of ones.
- f(A, B) = \sum_{(u,v)} |P_A - P_B| otherwise, where the sum is over all edges of the tree, P_A denotes the number of vertices with value 1 in A in one of the components after deleting edge (u, v), and P_B is the number of vertices with value 1 in B in the same component.
- To solve the entire problem, for each edge (u, v) we added\sum_{P_A = 0}^N \binom{y_1}{P_A-x_1} \cdot \binom{y_2}{C-P_A-x_2} \cdot |P_A-P_B|to the answer, where:
- x_1 and y_1 denote the number of ones and ?'s among values of A in the component containing u after deleting the edge.
- x_2 and y_2 denote the number of ones and ?'s among values of A in the component containing v.
- C denotes the total number of ones in string B.
- P_B denotes the number of ones in B in the component containing u.
- P_A denotes the target number of ones in A in the component containing u.
To optimize the above computation, observe that the structure of the tree doesn’t actually matter too much when computing the answer for an edge.
Specifically, once the edge is fixed, its contribution becomes purely a function of x_1, y_1, x_2, y_2, C,P_B.
In fact, x_2 and y_2 can be written in terms of x_1 and y_1 respectively (because x_1+x_2 equals the total number of ones in A, and y_1+y_2 equals the total number of ?'s in A.)
Further, C is a global constant.
So, the contribution of an edge is a function of only the values x_1, y_1, P_B.
In particular, if we let Y denote the count of ? in A, and X denote C - (\text{count of ones in A}), the expression we end up with is
Note that rather than sum over only P_A \in [0, N], we can pretend the sum is over all integers P_A and the value doesn’t change.
Let x = P_A - x_1 and d = P_B - x_1.
Then the above expression is equivalent to
which is visually a lot simpler to look at (and work with.)
Note here that X and Y are global constants, so the above expression is purely a function of y_1 and d.
Let’s denote this value as g(y_1, d).
We have N-1 queries for computing g with various inputs.
However, we don’t need to answer them all separately - we can try to answer them in some order that’s convenient to us.
In particular, suppose we know the value of g(y, d) already. Let’s try to compute g(y, d+1) from it.
The new expression is
which behaves differently for x \ge d+1 and x \lt d+1.
For the former, we want to subtract \binom{y}{x}\binom{Y-y}{X-x} while for the latter, we want to add this term.
To handle this, note that we’re adding all terms of the prefix till d and subtracting all other terms.
So, along with the value of g(y, d), we can also store the sum of all terms of the form \binom{y}{x}\binom{Y-y}{X-x}, and also the prefix sum of these terms with x \le d.
With these two values, computing the change from g(y, d) to g(y, d+1) is trivial in constant time; and the new prefix sum (with updated d) can also be maintained easily since it’s a difference of a single term.
Similarly, from g(y, d) we can move to g(y, d-1) easily.
Next, let’s try to move from g(y, d) to g(y+1, d).
The new expression is
With some algebraic manipulation, it can be shown that
Note that the final summation term can be once again written in terms of the prefix sums of \binom{y}{x}\binom{Y-y}{X-x} till d, along with the prefix sums of \binom{y}{x}\binom{Y-y}{X-x}x, also till d.
This gives us a way to compute g(y+1, d) from g(y, d) quickly, as long as this additional information is maintained.
And indeed, it is possible to store the prefix sums of \binom{y}{x}\binom{Y-y}{X-x} and \binom{y}{x}\binom{Y-y}{X-x}x up to d, and keep these two values updated as d and/or y change by 1.
(Doing it for d is trivial, doing it for y requires some algebraic manipulation which is left as an exercise to the reader - see the attached code for the exact transformation.)
Similarly, g(y-1, d) can be computed from g(y, d) quickly too.
Now that we’re able to move from (y, d) to adding/subtracting 1 to either one of its coordinates quickly, we can use Mo’s algorithm to answer all our queries.
The time complexity is \mathcal{O}(N\sqrt N\log {MOD}) because of the divisions involved.
If necessary, it’s possible to optimize out the log factor - we only perform divisions with values \le N and these inverses can all be precomputed.
TIME COMPLEXITY:
\mathcal{O}(N\sqrt N) 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());
/**
* 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);
auto fac = precalc_factorial<Zp>(200005);
auto ifac = fac;
for (auto &x : ifac) x = 1/x;
auto C = [&] (int n, int r) {
if (n < r or r < 0) return Zp(0);
return fac[n] * ifac[r] * ifac[n-r];
};
auto inv = fac;
for (int i = 1; i < 200005; ++i) inv[i] = Zp(1)/i;
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);
vector<array<int, 2>> queries;
for (int u = 1; u < n; ++u) {
auto [x1, y1, pb] = subct[u];
queries.push_back({y1, pb - x1});
}
Zp ans = 0;
auto solve = [&] () {
const int B = sqrt(n) + 2;
sort(begin(queries), end(queries), [&](auto &q1, auto &q2) {
auto [l1, r1] = q1;
auto [l2, r2] = q2;
if (l1/B != l2/B) return l1/B < l2/B;
if ((l1/B)%2 == 0) return r1 < r2;
else return r1 > r2;
});
int X = ranges::count(b, '1') - ranges::count(a, '1');
int Y = ranges::count(a, '?');
auto f = [&] (int y, int x) {
return C(y, x) * C(Y-y, X-x);
};
Zp val = 0;
Zp ps1 = 0, ps2 = 0;
// init with (0, 0)
for (int x = 0; x <= n; ++x) {
val += C(0, x) * C(Y, X-x) * abs(x);
}
int cury = 0, curd = 0;
auto addy = [&] () {
val += C(Y-1, X-1) - 2*(X*ps1 - ps2) * inv[(Y-cury)];
Zp tmp = f(cury, curd-1) * (X - curd + 1) * inv[(Y - cury)];
ps2 += (X * (ps1 - f(cury, curd-1)) - ps2 + f(cury, curd-1)*(curd-1)) * inv[(Y - cury)] - tmp * (curd-1);
ps1 -= tmp;
};
auto suby = [&] () {
ps1 += f(cury, curd) * curd * inv[cury];
ps2 -= ps2 * inv[cury] - (curd - 1)*f(cury, curd)*curd*inv[cury];
val -= C(Y-1, X-1) - 2*(X*ps1 - ps2)*inv[(Y-cury+1)];
};
auto addd = [&] () {
ps1 += f(cury, curd);
ps2 += f(cury, curd) * curd;
val += 2*ps1 - C(Y, X);
};
auto subd = [&] () {
val -= 2*ps1 - C(Y, X);
ps1 -= f(cury, curd-1);
ps2 -= f(cury, curd-1) * (curd-1);
};
for (auto [y, d] : queries) {
while (cury > y) {
suby();
--cury;
}
while (curd < d) {
addd();
++curd;
}
while (cury < y) {
addy();
++cury;
}
while (curd > d) {
subd();
--curd;
}
ans += val;
}
};
solve();
cout << ans << '\n';
}
}