What if Shaunak never uses the special move?

If Shaunak uses the special move, is using it on the first turn any different from using it later?

The above two cases are independent, and give you two conditions for Shaunak to win. What are they?

In a game without special moves, every turn reduces the total sum of the array, $$$\sum a_i$$$, by exactly $$$1$$$. Thus, the game lasts exactly $$$\sum a_i$$$ turns. The first player (Shaunak) wins if $$$\sum a_i$$$ is odd, and the second player (Yash) wins if $$$\sum a_i$$$ is even.

• If $$$\sum a_i$$$ is odd: Shaunak plays normally and wins. He doesn't need the special move.
• If $$$\sum a_i$$$ is even: Shaunak is guaranteed to lose normally, so he must use his special move. Using it on his first turn sets the new array sum to $$$n \cdot k$$$. Now, Yash is forced to play first in a standard game. Following the parity logic, Yash loses (and Shaunak wins) if $$$n \cdot k$$$ is even.

Therefore, Shaunak wins if $$$\sum a_i$$$ is odd or if $$$n \cdot k$$$ is even. Otherwise, Yash wins.

#include <iostream>
 
int main()
{
    int t = 1;
    std::cin >> t;
 
    while (t--) {
        int n, k;
        std::cin >> n >> k;
        bool ans = ((n * k) % 2 == 0); // Either n*k is even
 
        int tot = 0; 
        while (n--) {
            int x;
            std::cin >> x; 
            tot += x;
        } 
        ans |= tot & 1; // Or sum(a) is odd
        
        std::cout << (ans ? "YES" : "NO") << std::endl;
    }
 
    return 0;
}

t = int(input())
 
for _ in range(t):
    n, k = map(int, input().split())
    arr = list(map(int, input().split()))
    
    tot = sum(arr)
    ans = (tot % 2 == 1) or ((n * k) % 2 == 0)
    
    print("YES" if ans else "NO")

Think of the array in terms of segments instead of individual elements. What really matters is where the value changes. Define a boundary as a position where adjacent elements differ. How does a flip operation affect boundaries?

A valid operation must include the special index. So one endpoint of the flip lies to the left of it, and the other to the right (or at it). What does this imply about which boundaries you can remove together?

You can pair boundaries from the left and right. What happens to the remaining ones after pairing?

Let $$$V$$$ be the target value (the value at the special index $$$p_1$$$). Construct a binary array $$$b$$$ of size $$$n+2$$$ (0-indexed) where $$$b_i = 1$$$ if $$$a_i = V$$$, and $$$0$$$ otherwise. Pad the array with $$$b_0 = 1$$$ and $$$b_{n+1} = 1$$$.

Define a boundary as an index $$$j$$$ where $$$b_j \neq b_{j+1}$$$. Our goal is to eliminate all boundaries.

• Flipping a range $$$[l, r]$$$ toggles the boundaries at $$$l-1$$$ and $$$r$$$.
• The pivot constraint $$$l \le p_1 \le r$$$ means one toggled boundary must be $$$ \lt p_1$$$ and the other $$$\ge p_1$$$.

This naturally divides the boundaries into two independent regions:

1. Region 1: Boundaries in $$$[0, p_1 - 1]$$$. Let the count be $$$x_1$$$.
2. Region 2: Boundaries in $$$[p_1, n]$$$. Let the count be $$$x_2$$$.

Each valid operation removes at most one boundary from Region 1 and one from Region 2. We can pair boundaries across regions at a cost of 1 operation per pair. Once the smaller region is exhausted, the remaining boundaries in the larger region must be eliminated individually.

Thus, the minimum operations required are:

Time Complexity: $$$O(n)$$$

#include <bits/stdc++.h>
using namespace std;
 
void solve() {
    int n, k;
    cin >> n >> k;
    vector<int> arr(n+2);
    for (int i = 1; i <= n; i++) {
        cin >> arr[i];
    }
    int pivot;
    cin >> pivot;
 
    arr[0] = arr[n+1] = arr[pivot];
 
    int countL = 0, countR = 0;
    for (int i = 0; i < pivot; i++) {
        if (arr[i] != arr[i+1]) {
            countL++;
        }
    }
    for (int i = pivot; i < n+1; i++) {
        if (arr[i] != arr[i+1]) {
            countR++;
        }
    }
 
    cout << max(countL, countR) << "\n";
}
 
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    int t;
    cin>>t;
    while(t--) {
        solve();
    }
    return 0;
}

Think of a one-dimensional array that wraps around itself. Check whether repeatedly jumping by a fixed step on such an array can ever let you visit every position or just a subset.

Treat each position as two states: $$$(i,j,\text{last move was Down})$$$ and $$$(i,j,\text{last move was Right})$$$. Follow the sequence and count how many such states appear before the path cycles. Compare this with the total possible states, i.e., $$$2 \cdot n \cdot m$$$.

To visit every row and column, we strictly require $$$\gcd(n, a) = 1$$$ and $$$\gcd(m, b) = 1$$$. Assume these hold.

A full cycle of moves (1 Down, 1 Right) yields a displacement of $$$(+a, +b)$$$. To return to $$$A_{1,1}$$$ with the exact same next move, we need $$$p$$$ cycles where:

The minimal such $$$p$$$ is $$$\text{lcm}(n, m)$$$. Thus, the path loops entirely after $$$2 \cdot \text{lcm}(n, m)$$$ moves.

Because the maximum number of unique states Prakul can visit is $$$2 \cdot \text{lcm}(n, m)$$$, we compare this with the total number of grid tiles.

• If $$$\gcd(n, m) \ge 3$$$:
  The grid size is $$$n \cdot m = \gcd(n, m) \cdot \text{lcm}(n, m) \ge 3 \cdot \text{lcm}(n, m)$$$.
  This strictly exceeds the maximum visitable states ($$$2 \cdot \text{lcm}(n, m)$$$), so covering the grid is impossible.


• If $$$\gcd(n, m) = 1$$$:
  Here, $$$2 \cdot \text{lcm}(n, m) = 2 \cdot n \cdot m$$$.
  The loop length is twice the number of grid tiles. By symmetry, every tile is visited exactly twice (once in each direction). Hence, the grid is fully covered.


• If $$$\gcd(n, m) = 2$$$:
  Here, $$$2 \cdot \text{lcm}(n, m) = n \cdot m$$$.
  Prakul covers the grid if and only if he never visits a tile in both direction states. Let's assume a tile is visited in both states.
  
  In State 1 (ready to move Right), after $$$k_1$$$ full cycles, his displacement is $$$k_1 \cdot a$$$ rows and $$$k_1 \cdot b$$$ columns.
  In State 2 (ready to move Down), he has completed $$$k_2$$$ full cycles plus one extra Right move, making his displacement $$$k_2 \cdot a$$$ rows and $$$(k_2 + 1) \cdot b$$$ columns.
  
  For these to be the exact same tile in the wrapping grid, their displacements must be equivalent: $$$ k_1 \cdot a \equiv k_2 \cdot a \pmod n \Rightarrow k_1 \equiv k_2 \pmod n $$$ $$$ k_1 \cdot b \equiv (k_2 + 1)\cdot b \pmod m \Rightarrow k_1 \equiv k_2 + 1 \pmod m $$$ Letting the difference in cycles be $$$x = k_1 - k_2$$$. Then: $$$ x \equiv 0 \pmod n, \quad x \equiv 1 \pmod m $$$ This implies: $$$ c \cdot n - d \cdot m = 1 $$$ By Bézout's identity, this is possible only if $$$\gcd(n, m)$$$ divides $$$1$$$. But $$$\gcd(n, m) = 2$$$, so this is impossible. Hence, no tile is visited in both states, and exactly $$$n \cdot m$$$ distinct tiles are visited.

Therefore, Prakul covers the grid if and only if $$$\gcd(n, a) = 1$$$, $$$\gcd(m, b) = 1$$$, and $$$\gcd(n, m) \le 2$$$.

#include <iostream>
#include <numeric>
 
int main()
{
    int t = 1;
    std::cin >> t;
 
    while (t--) {
        long long n, m, a, b;
        std::cin >> n >> m >> a >> b;
        bool possible = (
            std::gcd(n, a) == 1 and
            std::gcd(m, b) == 1 and
            std::gcd(n, m) <= 2
        );
        std::cout << (possible ? "YES" : "NO") << std::endl;
    }
 
    return 0;
}

Similar to the easy version, think in terms of the boundaries where adjacent values differ. But now, the special indices divide the array into multiple regions.

You want to maximize how many times you can pair two boundaries from different regions (removing two in one operation). When is this always possible?

Let $$$S$$$ be the total number of boundaries and $$$X$$$ the maximum number in any region. When does one region have too many boundaries compared to the rest?

If one region dominates (has more than half the total boundaries), some boundaries are forced to be removed individually. Otherwise, all boundaries can be paired.

Similar to the easy version, construct a padded binary array $$$b$$$ where $$$b_i = 1$$$ if $$$a_i$$$ matches the target value, else $$$0$$$. A boundary exists at $$$j$$$ if $$$b_j \neq b_{j+1}$$$. Flipping $$$[l, r]$$$ toggles boundaries at $$$l-1$$$ and $$$r$$$.

The constraint that $$$[l, r]$$$ must contain at least one pivot $$$p_i$$$ means we can only pair two boundaries in a single operation if they are separated by at least one pivot.

The $$$k$$$ pivots divide the boundaries into $$$k + 1$$$ independent regions:

• Region 1: $$$[0, p_1 - 1]$$$
• Region $$$i$$$: $$$[p_{i-1}, p_i - 1]$$$ for $$$2 \le i \le k$$$
• Region $$$k+1$$$: $$$[p_k, n]$$$

Let $$$S$$$ be the total number of boundaries across all regions, and $$$X$$$ be the maximum number of boundaries in any single region.

We want to maximize pairing boundaries across different regions (which removes 2 boundaries for 1 operation). Removing a boundary stuck in its own region costs 1 operation per boundary.

1. Case 1 ($$$X \le S/2$$$): No region dominates. We can optimally pair every boundary with one from a different region. Total operations: $$$S / 2$$$.
2. Case 2 ($$$X \gt S/2$$$): One region has more boundaries than the rest combined. We pair all $$$S - X$$$ "outsider" boundaries with boundaries from the majority region. The remaining $$$X - (S - X) = 2X - S$$$ boundaries in the majority region must be resolved individually. Total operations: $$$(S - X) + (2X - S) = X$$$.

Combining these conditions gives the minimum number of operations:

Time complexity: $$$O(n)$$$

#include <bits/stdc++.h>
using namespace std;
 
void solve() {
    int n, k;
    cin >> n >> k;
    vector<int> arr(n+2);
    for (int i = 1; i <= n; i++) {
        cin >> arr[i];
    }
    vector<int> pivots(k+2);
    for (int i = 1; i <= k; i++) {
        cin >> pivots[i];
    }
    arr[0] = arr[n+1] = arr[pivots[1]];
    pivots[k+1] = n+1;
 
    int S = 0, X = 0;
    for (int i = 0; i <= k; i++) {
        int count = 0;
        for (int j = pivots[i]; j < pivots[i+1]; j++) {
            if (arr[j] != arr[j+1]) {
                count++;
            }
        }
        S += count;
        X = max(X, count);
    }
 
    cout << max(S / 2, X) << "\n";
}
 
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    int t;
    cin>>t;
    while(t--) {
        solve();
    }
    return 0;
}

For each index $$$i$$$, the only indices that can ever contribute to $$$d_i$$$ are those with $$$j \gt i$$$ and $$$p_j \gt p_i$$$; so first check whether $$$d_i$$$ exceeds their count.

What happens if you process indices in decreasing order of $$$p_i$$$? Then one part of the domination condition is already automatically satisfied.

There are a few ways to approach this problem, but the most elegant solutions involve greedy construction. We will discuss two distinct but equally valid greedy strategies.

Method 1: Insertion by Decreasing $$$p_i$$$
This method determines the relative order of the elements in $$$q$$$ by processing the array based on the values of $$$p_i$$$.

Let's process the indices $$$i$$$ in decreasing order of their $$$p_i$$$ values (from $$$n$$$ down to $$$1$$$). We will maintain a list, ord, of the already processed indices. This list represents the elements sorted by their future $$$q$$$ values.

When we process the current index $$$i$$$, every index $$$j$$$ already in ord satisfies $$$p_j \gt p_i$$$. This perfectly handles the second condition of domination. Now we just need to satisfy the remaining conditions: $$$j \gt i$$$ and $$$q_j \gt q_i$$$.

Let $$$m$$$ be the number of indices $$$j$$$ currently in ord such that $$$j \gt i$$$. These are the only candidates that could possibly dominate $$$i$$$.

• If $$$d_i \gt m$$$, it is impossible to satisfy the condition, and we should output -1.
• Otherwise, to ensure exactly $$$d_i$$$ elements dominate $$$i$$$, we need exactly $$$d_i$$$ of these valid $$$j$$$ indices to have a larger $$$q$$$ value (meaning they are placed after $$$i$$$ in ord).
• Consequently, exactly $$$m - d_i$$$ of these valid $$$j$$$ indices must be placed before $$$i$$$ in ord.

We can simply iterate through ord, count until we have passed $$$m - d_i$$$ elements that satisfy $$$j \gt i$$$, and insert $$$i$$$ immediately after them. Once all $$$n$$$ indices are inserted into ord, their 1-based positions in the list become their actual $$$q_i$$$ values.

Method 2: Assigning $$$q_i$$$ from Large to Small
Alternatively, we can build $$$q$$$ directly by deciding which index receives the value $$$n$$$, then $$$n-1$$$, down to $$$1$$$.

Suppose we are trying to place the largest available value. The index $$$i$$$ that receives this value will have $$$q_i$$$ greater than all remaining unassigned elements. Therefore, nothing processed after $$$i$$$ can possibly dominate it. For $$$i$$$ to be a valid choice, its current $$$d_i$$$ must be exactly $$$0$$$.

If there are multiple indices with $$$d_i = 0$$$, which one should we pick? We should greedily pick the one with the smallest index $$$i$$$. Picking a smaller index avoids "blocking" elements to its left that might need to be dominated by it later.

Once we select this index $$$i$$$ and assign it the current maximum value, it becomes a valid dominator for any unassigned index $$$j$$$ that sits to its left ($$$j \lt i$$$) and has a smaller $$$p$$$ value ($$$p_j \lt p_i$$$). We simulate this by decrementing $$$d_j$$$ by $$$1$$$ for all such valid $$$j$$$. We repeat this process until all elements are assigned. If at any step no unassigned element has $$$d_i = 0$$$, the array is invalid, and we output -1.

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int t;
    cin >> t; 

    while (t--) {
        int N;
        cin >> N;

        vector<int> P(N + 1), pos(N + 1);
        for (int i = 1; i <= N; i++) {
            cin >> P[i];
            pos[P[i]] = i; 
        }

        vector<int> D(N + 1);
        for (int i = 1; i <= N; i++) cin >> D[i];

        vector<int> ord(N);

        bool ok = true;

        for (int val = N; val >= 1 && ok; --val) {
            int i = pos[val];

            int m = 0;
            for (int idx : ord) if (idx > i) ++m;

            if (D[i] > m) {
                ok = false;
                break;
            }

            int needBefore = m - D[i]; 
            int cnt = 0;
            int insertPos = 0;

            while (insertPos < (int)ord.size() && cnt < needBefore) {
                if (ord[insertPos] > i) ++cnt;
                ++insertPos;
            }

            ord.insert(ord.begin() + insertPos, i);
        }

        if (!ok) {
            cout << -1 << '\n';
            continue;
        }

        vector<int> A(N + 1, 0);
        for (int k = 0; k < N; k++) {
            A[ord[k]] = k + 1;
        }

        for (int i = 1; i <= N; i++) {
            cout << A[i] << (i == N ? '\n' : ' ');
        }
    }

    return 0;
}

Model the game as a Nim game. Each interval contributes two independent piles. The whole game becomes a 4-pile Nim game. Who wins depends only on the XOR (Nim-sum) of all piles.

Alice can control the XOR contribution of the first interval. What values of XOR can she achieve?

Let $$$X$$$ be Alice's XOR and $$$Z$$$ be the XOR from the random interval. Alice loses only when $$$X = Z$$$. So instead of maximizing wins directly, minimize how often $$$Z = X$$$ happens.

Fix a value of $$$X$$$. Count how many pairs $$$(a, b)$$$ (derived from the second interval) satisfy $$$a \oplus b = X$$$ under the constraint $$$a + b \le x_2 - 1$$$.

Use the identity:

This reduces the problem to counting integers with a bitmask constraint.

We can model this problem using the standard Game of Nim. Consider a single interval $$$[l, r]$$$ constrained within $$$[1, x]$$$.

• Decreasing $$$l$$$ (shifting left boundary) is equivalent to reducing a pile of size $$$l-1$$$.
• Increasing $$$r$$$ (shifting right boundary) is equivalent to reducing a pile of size $$$x-r$$$.

Thus, the game is played on 4 piles of stones with sizes:

The game ends when no moves are possible (all piles are size 0). Alice wins if the Nim-sum (XOR sum) of all pile sizes is non-zero.

Let $$$X$$$ be the XOR sum of the first interval's piles:

Let $$$Z$$$ be the XOR sum of the second interval's piles:

Alice wins if $$$X \oplus Z \neq 0$$$, which implies $$$Z \neq X$$$. Since the second interval is chosen uniformly at random, to maximize her winning probability, Alice must choose $$$l_1, r_1$$$ (and thus determine $$$X$$$) such that the number of outcomes where $$$Z = X$$$ is minimized.

Alice can produce any XOR value $$$X$$$ in the range $$$[0, x_1 - 1]$$$. A simple construction to achieve a specific XOR value $$$X$$$ is to set the first pile to 0 and the second to $$$X$$$.

• Set $$$l_1 = 1$$$ (Pile size $$$1-1=0$$$)
• Set $$$r_1 = x_1 - X$$$ (Pile size $$$x_1 - (x_1 - X) = X$$$)

This gives an interval $$$[1, x_1 - X]$$$, which is valid for any $$$X \lt x_1$$$.

We need to count how many valid intervals $$$[l_2, r_2]$$$ produce an XOR sum equal to Alice's chosen $$$X$$$. Let the pile sizes for the second interval be $$$a = l_2 - 1$$$ and $$$b = x_2 - r_2$$$. The constraints on the second interval are:

We are looking for the count of pairs $$$(a, b)$$$ such that $$$a \oplus b = X$$$ and their sum fits in the range. We use the fundamental relationship between sum and XOR:

Let $$$Y = a + b$$$. Substituting $$$a \oplus b = X$$$, we get:

Rearranging for the AND term:

This equation gives us the necessary conditions for a valid sum $$$Y$$$:

1. $$$Y \ge X$$$ (since $$$a \ \& \ b \ge 0$$$)
2. $$$Y - X$$$ must be divisible by 2
3. Crucial Condition: The bits set in $$$(a \ \& \ b)$$$ cannot be set in $$$(a \oplus b)$$$. Mathematically: $$$((Y - X) / 2) \ \& \ X = 0$$$

If these conditions are met, how many pairs $$$(a, b)$$$ sum to $$$Y$$$ with XOR $$$X$$$? For every bit set in $$$X$$$, exactly one of $$$a$$$ or $$$b$$$ must have that bit set (2 choices: $$$1/0$$$ or $$$0/1$$$). The bits in the Bitwise-AND of $$$a$$$ and $$$b$$$ are fixed to 1 for both. Thus, for a valid $$$Y$$$, there are $$$2^{\text{popcount}(X)}$$$ ways to form the pairs.

We can simplify the condition by letting $$$K = (Y - X) / 2$$$. The condition $$$Y \le x_2 - 1$$$ becomes:

Let $$$Limit = \left\lfloor \frac{x_2 - 1 - X}{2} \right\rfloor$$$. The problem reduces to finding the count of integers $$$K$$$ such that:

1. $$$0 \le K \le Limit$$$
2. $$$K \ \& \ X = 0$$$

This specific counting problem (finding numbers less than a limit with a mask constraint) can be solved using simple bitwise logic (Greedy or Digit DP style) in $$$O(\log(\text{limit}))$$$.

Therefore:

• Iterate $$$X$$$ from $$$0$$$ to $$$x_1 - 1$$$
• Calculate $$$Limit = (x_2 - 1 - X) / 2$$$
• Count valid $$$K$$$'s in $$$[0, Limit]$$$ disjoint from $$$X$$$
• Total configurations for this $$$X$$$ is: $$$(\text{Count of } K) \times 2^{\text{popcount}(X)}$$$
• Pick the $$$X$$$ that minimizes this total

There are $$$x_1$$$ choices for $$$X$$$. The counting step takes $$$O(\log x_2)$$$. Total Complexity: $$$O(x_1 \log x_2)$$$.

#include <bits/stdc++.h>
using namespace std;
 
long long count_valid_k(int limit, int mask) {
    if (limit < 0) return 0;
 
    long long count = 0;
 
    for (int i = 20; i >= 0; i--) {
        int bit_limit = (limit >> i) & 1;
        int bit_mask = (mask >> i) & 1;
 
        int lower_free_bits = 0;
        for (int j = 0; j < i; j++) {
            if (!((mask >> j) & 1)) {
                lower_free_bits++;
            }
        }
 
        if (bit_mask == 1) {
            if (bit_limit == 1) {
                count += (1LL << lower_free_bits);
                return count;
            }
        } else {
            if (bit_limit == 1) {
                count += (1LL << lower_free_bits);
            }
        }
    }
 
    return count + 1;
}
 
void solve() {
    int x1, x2;
    cin >> x1 >> x2;
 
    long long min_matching_outcomes = -1;
    int best_X = -1;
 
    for (int X = 0; X < x1; X++) {
        int val = x2 - 1 - X;
        long long current_matching_outcomes = 0;
 
        if (val < 0) {
            current_matching_outcomes = 0;
        } else {
            int limit = val / 2;
            long long k_count = count_valid_k(limit, X);
            current_matching_outcomes = k_count * (1LL << __builtin_popcount(X));
        }
 
        if (min_matching_outcomes == -1 || current_matching_outcomes < min_matching_outcomes) {
            min_matching_outcomes = current_matching_outcomes;
            best_X = X;
        }
 
        if (min_matching_outcomes == 0) break;
    }
 
    int l1 = 1;
    int r1 = x1 - best_X;
 
    cout << l1 << " " << r1 << "\n";
}
 
int main() {
    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}

If we consider only flipping on root-to-$$$u$$$ paths, how many labellings are possible such that the cost is exactly $$$k$$$?

Answer: There will be exactly $$$\binom{n}{k}$$$ labellings for a tree of size $$$n$$$, regardless of the tree's shape.

A path through the root $$$r$$$ consists of a root-to-node path in the left subtree $$$L$$$, the root $$$r$$$ itself, and a root-to-node path in the right subtree $$$R$$$ (where either subtree path can be empty).

Because we can pair one needed flip in $$$L$$$ with one needed flip in $$$R$$$ into a single operation, clearing both subtrees requires exactly $$$m = \max(S(L), S(R))$$$ operations.

These $$$m$$$ operations will flip the root's label exactly $$$m$$$ times. If the root's final label is still $$$1$$$, we require exactly $$$1$$$ additional operation (flipping just the root itself, using empty paths for $$$L$$$ and $$$R$$$).

Thus, the total cost to clear the tree is:

For a fixed pair of subtrees (which fixes $$$m$$$), we have two choices for the root's initial label $$$a(r) \in {0, 1}$$$. Exactly one choice results in a total cost of $$$m$$$, and the other results in a cost of $$$m + 1$$$.

Consequently, the tree will have a cost of exactly $$$k$$$ if and only if $$$m \in {k, k-1}$$$. This is mathematically equivalent to the condition $$$m \le k$$$ minus the condition $$$m \le k - 2$$$.

For a fixed split of sizes $$$i$$$ and $$$j = n - 1 - i$$$, what is the number of valid $$$(L, R)$$$ pairs such that $$$\max(S(L), S(R)) \le t$$$? Calculate this, then iterate $$$i$$$ from $$$0$$$ to $$$n-1$$$.

The problem can be drastically simplified by transforming the tree labels into a different state. Let $$$a(v) \in {0, 1}$$$ be the initial label of node $$$v$$$. For any node $$$v$$$, define a new value:

(If a child does not exist, we treat its label as $$$0$$$).

Why is this transformation useful? Consider what happens to $$$x(v)$$$ when we perform an operation on a simple path from the root to some node $$$u$$$: * If $$$v$$$ is on the path but is not the endpoint $$$u$$$, both $$$a(v)$$$ and exactly one of its children's labels are flipped. Thus, $$$x(v)$$$ changes by $$$1 \oplus 1 = 0$$$. * If $$$v$$$ is the endpoint $$$u$$$, $$$a(v)$$$ is flipped, but neither of its children is on the path. Thus, $$$x(v)$$$ changes by $$$1 \oplus 0 = 1$$$. * If $$$v$$$ is not on the path, neither $$$a(v)$$$ nor its children are flipped. Thus, $$$x(v)$$$ changes by $$$0$$$.

An operation on a root-to-$$$u$$$ path toggles exactly $$$x(u)$$$ and leaves all other $$$x(v)$$$ values unchanged!

Furthermore, the mapping from $$$a(v)$$$ to $$$x(v)$$$ is a perfect bijection for any fixed tree structure (you can uniquely reconstruct $$$a(v)$$$ from $$$x(v)$$$ working bottom-up). Therefore, to zero-out a subtree $$$T$$$ using only paths starting from its root, we must flip exactly the paths ending at nodes where $$$x(v) = 1$$$. The minimum number of operations for a subtree $$$T$$$ is simply $$$S(T) = \sum_{v \in T} x(v)$$$. Since $$$x(v) \in {0, 1}$$$, the number of labelings of a tree of size $$$m$$$ that yield exactly $$$S(T) = s$$$ is $$$\binom{m}{s}$$$, independent of the tree's shape.

Evaluating the Whole Tree: An allowed operation is a path through the root $$$r$$$. This path consists of a root-to-node path in the left subtree $$$L$$$, the root $$$r$$$ itself, and a root-to-node path in the right subtree $$$R$$$ (where either subtree path can be empty).

Because we can pair one needed flip in $$$L$$$ with one needed flip in $$$R$$$ into a single operation, clearing both subtrees requires exactly $$$m = \max(S(L), S(R))$$$ operations.

These $$$m$$$ operations will flip the root's label exactly $$$m$$$ times. If the root's final label is still $$$1$$$, we require exactly $$$1$$$ additional operation (flipping just the root itself, using empty paths for $$$L$$$ and $$$R$$$). Thus, the total cost to clear the tree is:

For a fixed pair of subtrees (which fixes $$$m$$$), we have two choices for the root's initial label $$$a(r) \in {0, 1}$$$. Exactly one choice results in a total cost of $$$m$$$, and the other results in a cost of $$$m + 1$$$. Consequently, the tree will have a cost of exactly $$$k$$$ if and only if $$$m \in {k, k-1}$$$. This is mathematically equivalent to the condition $$$m \le k$$$ minus the condition $$$m \le k - 2$$$.

Counting the Trees: Let $$$C_t$$$ be the $$$t$$$-th Catalan number (the number of binary tree structures on $$$t$$$ nodes). Define the prefix sum of binomial coefficients as:

(with $$$BS(m, t) = 0$$$ for $$$t \lt 0$$$).

For a fixed split of sizes $$$i$$$ and $$$j = n - 1 - i$$$, the number of valid $$$(L, R)$$$ pairs such that $$$\max(S(L), S(R)) \le t$$$ is given by:

To find the number of trees with cost exactly $$$k$$$, we subtract the valid pairs for $$$t = k-2$$$ from the valid pairs for $$$t = k$$$:

Complexity and Optimization: We need to calculate $$$BS(m, t)$$$ for all $$$m$$$ up to $$$n$$$, for fixed values of $$$t \in {k, k-2}$$$. Recomputing the sum naively would be too slow. Instead, we can use the identity:

This allows us to compute $$$BS(m, k)$$$ and $$$BS(m, k-2)$$$ iteratively in $$$\mathcal{O}(1)$$$ transitions. The Catalan numbers can also be precomputed in linear time. Thus, the entire summation takes $$$\mathcal{O}(n)$$$ time per test case.

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
 
using namespace std;
#define int long long
#define all(s) s.begin(), s.end()
#define sz(s) (int)s.size()
#define testcases \
    cin >> tt;    \
    for (i = 1; i <= tt; i++)
#define fast                      \
    ios_base::sync_with_stdio(0); \
    cin.tie(NULL);                \
    cout.tie(NULL)
 
const int M = 1e9 + 7;
 
using i32 = int32_t;
template <const auto mod>
struct mint {
    template<typename T = i32>
    constexpr mint(T x = 0) : val(x % mod + (x < 0) * mod) {}
    mint &operator+=(const mint &b) {
        val += b.val;
        val -= mod * (val >= mod);
        return *this;
    }
    mint &operator-=(const mint &b) {
        val -= b.val;
        val += mod * (val < 0);
        return *this;
    }
    mint &operator*=(const mint &b) {
        val = 1ll * val * b.val % mod;
        return *this;
    }
    mint &operator/=(const mint &b) { return *this *= b.inv(); }
    mint inv() const {
        i32 x = 1, y = 0, t;
        for (i32 a = val, b = mod; b; swap(a, b), swap(x, y))
            t = a / b, a -= t * b, x -= t * y;
        return mint(x);
    }
    mint pow(int b) const {
        mint a = *this, res(1);
        for (; b; a *= a, b /= 2)
            if (b & 1) res *= a;
        return res;
    }
    friend mint operator+(const mint &a, const mint &b) { return mint(a) += b; }
    friend mint operator-(const mint &a, const mint &b) { return mint(a) -= b; }
    friend mint operator*(const mint &a, const mint &b) { return mint(a) *= b; }
    friend mint operator/(const mint &a, const mint &b) { return mint(a) /= b; }
    friend bool operator==(const mint &a, const mint &b) { return a.val == b.val; }
    friend bool operator!=(const mint &a, const mint &b) { return a.val != b.val; }
    friend bool operator<(const mint &a, const mint &b) { return a.val < b.val; }
    friend ostream &operator<<(ostream &os, const mint &a) { return os << a.val; }
    friend string to_string(const mint& a) { return to_string(a.val); } 
    i32 val;
};
using Mint = mint<M>;
typedef vector<Mint> vm;
typedef vector<vm> vvm;
 
#define PREC_TILL (int) 2'000'000
vector<Mint> factorials(PREC_TILL + 1);
vector<Mint> invFactorials(PREC_TILL + 1);
void prec() {
    factorials[0] = 1;
    for (int i = 1; i <= PREC_TILL; i++) factorials[i] = factorials[i - 1] * i;
    invFactorials.back() = factorials.back().inv();
    for (int i = PREC_TILL - 1; i >= 0; i--) {
        invFactorials[i] = invFactorials[i + 1] * (i + 1);
    }
}
 
Mint nCr(int n, int r) {
    if (r < 0 or r > n) return 0;
    return factorials[n] * invFactorials[n - r] * invFactorials[r];
}
 
Mint catalan(int n) {
    if (n < 0) return Mint(0);
    return factorials[2 * n] * invFactorials[n] * invFactorials[n + 1];    
}
 
void solve(__attribute__((unused)) int tt) {
    int n, k; cin >> n >> k;
    
    Mint ans = 0;
    
    for (auto y : {k - 1, k}) {
        vm prefbinom(n + 1);
        prefbinom[0] = 1;
        for (int i = 1; i <= n; i++) {
            prefbinom[i] = 2 * prefbinom[i - 1] - nCr(i - 1, y);
        }
 
        for (int i = 0, j = n - 1; i <= n - 1; i++, j--) {
            ans += 2 * catalan(i) * catalan(j) * nCr(i, y) * (prefbinom[j] - nCr(j, y));
        }
    
        for (int i = 0, j = n - 1; i <= n - 1; i++, j--) {
            ans += catalan(i) * catalan(j) * nCr(i, y) * nCr(j, y);
        }
    }
 
    cout << ans << "\n";
}
 
int32_t main() {
    fast;
    prec();
    int tt = 1;
    int i = 1;
    testcases
        solve(i);
}

The key observation is that every chosen edge belongs to a matching, so every vertex can participate in at most one swap. If the two copies of some badge are initially at distance at least $$$4$$$, then even after moving both of them by one edge, they still cannot become adjacent. Any deal that can ever be sealed must come from a very small neighbourhood. This suggests a tree DP.

Root the tree at vertex $$$1$$$. For a node $$$u$$$, after all swaps, the badge finally sitting on $$$u$$$ can only have come from:

• $$$u$$$ itself,
• $$$\mathrm{par}(u)$$$,
• or one of the children of $$$u$$$.

For every such vertex $$$x$$$, let $$$\text{f}[u][x]$$$ be the maximum total profit contributed by sealed deals whose upper endpoint lies inside the subtree of $$$u$$$, under the condition that after the reshuffle, vertex $$$u$$$ contains the badge that originally stood on vertex $$$x$$$.

Now define $$$\text{best}[u] = \max_{x \ne \mathrm{par}(u)} \text{f}[u][x]$$$, and $$$\text{sum}[u] = \sum_{v \text{ child of } u} \text{best}[v].$$$

Here, $$$\text{sum}[u]$$$ is the value obtained if every child subtree behaves independently in its best possible way.

Let's compute some state $$$\text{f}[u][x]$$$. Let $$$c = a_x$$$, i.e. the badge that ends up at $$$u$$$.

First, we account for how this badge arrives at $$$u$$$. If $$$x=u$$$ or $$$x=\mathrm{par}(u)$$$, no child subtree is forced yet. If $$$x$$$ is a child $$$v$$$ of $$$u$$$, then edge $$$(u,v)$$$ must be chosen in the matching. In that case, subtree $$$v$$$ cannot simply contribute $$$\text{best}[v]$$$; instead, after the swap, vertex $$$v$$$ must receive the old badge of $$$u$$$, so we must use the state $$$\text{f}[v][u]$$$. So compared to the default value $$$\text{sum}[u]$$$, bringing the badge from a child $$$v$$$ changes the contribution of that child by $$$\text{f}[v][u] - \text{best}[v]$$$.

Now we ask whether the deal with badge $$$c$$$ can be sealed at $$$u$$$. Let $$$y$$$ be the other occurrence of badge $$$c$$$.

Since one copy of $$$c$$$ already ends at $$$u$$$, the only way to count this deal at $$$u$$$ is for the other copy to end at some child of $$$u$$$. Because each badge moves by at most one edge, only two cases are possible.

The first case is that $$$y$$$ is already a child $$$t$$$ of $$$u$$$. Then we can seal badge $$$c$$$ on edge $$$(u,t)$$$. Inside subtree $$$t$$$, vertex $$$t$$$ must keep its own badge, so we replace $$$\text{best}[t]$$$ by $$$\text{f}[t][t]$$$, and gain $$$w_c$$$. The extra profit from this is $$$\text{f}[t][t] - \text{best}[t] + w_c$$$.

The second case is that $$$y$$$ is a grandchild of $$$u$$$: say $$$y$$$ is a child of $$$t$$$ and $$$t$$$ is a child of $$$u$$$. Then we can move the badge from $$$y$$$ up to $$$t$$$, and seal badge $$$c$$$ on edge $$$(u,t)$$$. In this case we replace $$$\text{best}[t]$$$ by $$$\text{f}[t][y]$$$, so the extra profit is $$$\text{f}[t][y] - \text{best}[t] + w_c$$$.

There is one important restriction here: this second case is invalid if $$$t=x$$$, because then $$$t$$$ would need to participate in two swaps, which is impossible since the chosen edges must form a matching.

If neither of the above cases applies, then badge $$$c$$$ cannot form a sealed deal at $$$u$$$.

Therefore, each transition is very small:

• start with $$$\text{sum}[u]$$$,
• maybe modify one child because its badge is sent to $$$u$$$,
• maybe modify one child because it helps seal the deal of the badge now at $$$u$$$,
• take the second modification only if it is profitable.

So the transition is essentially: $$$\text{f}[u][x] = \text{sum}[u] + \text{arrival adjustment} + \max(0,\text{sealing bonus})$$$.

That's the whole idea. Once we fix which badge ends up at $$$u$$$, all child subtrees are independent except for at most two of them.

We process the tree bottom-up. After computing all states for $$$u$$$, we set $$$\text{best}[u] = \max_{x \ne \mathrm{par}(u)} \text{f}[u][x]$$$. The answer for the whole test case is simply $$$\text{best}[1]$$$.

The number of states is linear: for each node, we only keep states corresponding to the node itself, its parent, and its children. Each transition is $$$O(1)$$$. With the map-based implementation in the reference code, the total complexity is $$$O(N \log N)$$$ per test case. If you want it enough, you can make it $$$O(N)$$$ :)

Original code from _istil, cleaned up for reference.

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

constexpr int MAXN = 200005;

vector<int> g[MAXN];
int a[MAXN], w[MAXN], parent_[MAXN], pos[MAXN][2];
map<int, ll> dp[MAXN];
ll sum[MAXN], best[MAXN];

// Stores the parent of each node in the rooted tree.
void dfs_parent(int u, int p) {
    parent_[u] = p;
    for (int v : g[u]) {
        if (v != p) {
            dfs_parent(v, u);
        }
    }
}

// Computes dp[u][from]:
// the badge currently at u originally came from vertex 'from'.
void update_state(int u, int color, int from, ll delta) {
    int other = pos[color][0] ^ pos[color][1] ^ from;

    if (parent_[other] == u) {
        dp[u][from] = sum[u] + delta + max(dp[other][other] - best[other] + w[color], 0LL);
    } else if (parent_[parent_[other]] == u && parent_[other] != from) {
        int mid = parent_[other];
        dp[u][from] = sum[u] + delta + max(dp[mid][other] - best[mid] + w[color], 0LL);
    } else {
        dp[u][from] = sum[u] + delta;
    }
}

void dfs_dp(int u) {
    sum[u] = 0;

    for (int v : g[u]) {
        if (v == parent_[u]) continue;
        dfs_dp(v);
        sum[u] += best[v];
    }

    // Badge at u stays at u.
    update_state(u, a[u], u, 0);

    // Badge at parent(u) moves to u.
    if (parent_[u] != 0) {
        update_state(u, a[parent_[u]], parent_[u], 0);
    }

    // Badge at child v moves to u.
    for (int v : g[u]) {
        if (v == parent_[u]) continue;
        update_state(u, a[v], v, dp[v][u] - best[v]);
    }

    best[u] = 0;
    for (auto [from, value] : dp[u]) {
        if (from != parent_[u]) {
            best[u] = max(best[u], value);
        }
    }
}

void solve() {
    int n;
    cin >> n;
    int N = 2 * n;

    for (int i = 1; i <= n; ++i) {
        cin >> w[i];
        pos[i][0] = pos[i][1] = 0;
    }

    for (int i = 1; i <= N; ++i) {
        cin >> a[i];
        pos[a[i]][pos[a[i]][0] == 0 ? 0 : 1] = i;
    }

    for (int i = 1; i <= N; ++i) {
        g[i].clear();
        dp[i].clear();
    }

    for (int i = 1; i < N; ++i) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    dfs_parent(1, 0);
    dfs_dp(1);

    cout << best[1] << '\n';
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int T;
    cin >> T;
    while (T--) {
        solve();
    }
    return 0;
}