You have to output the sum of the happiness over all possible combinations of choices.

Use the trick. At the end of the process, for each child color one of his cookies in red. You have to output the number of combination of choices, if you assume that the coloring is a choice too. What if you distribute red cookies during the $$$k$$$ days instead of coloring cookies at the end?

Now the problem becomes "In the $$$i$$$-th day, we will choose $$$a_i$$$ children and give either a red cookie or a normal cookie to each child chosen. In how many ways each children ends up having exactly one red cookie?"

Let's find a DP. You can simulate rounds. Which information do you need after each round?

After each round, you only need the number of children who have already received a red cookie.

If you use the trick, the problem becomes "In the $$$i$$$-th day, we will choose $$$a_i$$$ children and give either a red cookie or a normal cookie to each child chosen. In how many ways each children ends up having exactly one red cookie?".

You can calculate the answer with DP. Let $$$dp_{i, j}$$$ be the number of ways to distribute the candies in the first $$$i$$$ rounds and give a red cookie to exactly $$$j$$$ children. When you calculate $$$dp_{i, j}$$$, you can iterate over the number $$$m$$$ of red cookies distributed in the last round. The transition is $$$dp_{i, j} = \sum_{m=0}^{\min(j, a_i)} dp_{i-1, j-m} \cdot \binom{n-j+m}{m} \cdot \binom{n-m}{a_i-m}$$$. In fact, there are

• $$$dp_{i-1, j-m}$$$ ways to distribute candies in the previous rounds
• $$$\binom{n-j+m}{m}$$$ ways to choose the children who receive a red cookie
• $$$\binom{n-m}{a_i-m}$$$ ways to choose the children who receive a normal cookie

The answer is $$$dp_{k, n}$$$.

Complexity: $$$O(n^2k)$$$.

This problem is very similar to the previous one. The differences are that the boxes already contain some balls and $$$k$$$ is large (but the operation are simpler, because you insert exactly $$$1$$$ ball for each operation).

Let's deal with the first difference. In the previous problem, when did we use the fact that the boxes were initially empty?

Hint 5 in the previous problem is true because the boxes were initially empty (so they also were initially equal). In this problem, is it possible to avoid dealing with the $$$a_i$$$ during the operations (e.g. can you do some preprocessing before considering the operations)?

Let's color the balls already in the boxes. You can calculate the number of ways to color $$$m$$$ balls with DP.

Now we can iterate over the number $$$m$$$ of balls colored initially.

For each $$$m$$$, the problem becomes "In how many ways can you put $$$k$$$ balls in the boxes and color $$$n-m$$$ of them, one for each box without already colored balls?"

Instead of coloring the balls at the end, you can choose the operations where you put a red ball in a box. In how many ways can you do that?

Use the trick. First, fix the number of colored balls that were already in the boxes before the operations. Let $$$dp_{i, j}$$$ be the number of ways to color exactly $$$j$$$ balls in the first $$$i$$$ boxes. The transition is $$$dp_{i, j} = dp_{i-1, j} + a_i \cdot dp_{i-1, j-1}$$$.

If you color exactly $$$m$$$ boxes before the operations, the number of ways is $$$dp_{n, m}$$$, multiplied by the number of ways to put $$$k$$$ balls in the boxes and color $$$n-m$$$ of them, one for each box without already colored balls.

Now let's deal with the $$$k$$$ operations. Let's fix which operations add a red ball ($$$\binom{k}{n-m}$$$ ways), the order of red balls ($$$(n-m)!$$$ ways) and the boxes chosen in the other operations ($$$n^{k-n+m}$$$).

Therefore, the answer is $$$\frac{1}{n^m} \cdot \sum_{m=0}^{n} dp_{n, m} \cdot \binom{k}{n-m} \cdot (n-m)! \cdot n^{k-n+m}$$$

Complexity: $$$O(n^2)$$$.

You can replace the DP with a small to large NTT, and the complexity becomes $$$O(n \log^2 n)$$$.

In the official editorial, there is a solution with polynomials and linearity of expectation. Find the similarities between this solution and the one in the editorial.

$$$S$$$ can be obtained by only considering sequences where at least a person hands $$$0$$$ balls to the neighbor on his right.

Use the trick. For each person $$$i$$$, his red ball was initially owned either by $$$i$$$ or by the neighbor on his left.

Use DP. In how many ways can you distribute and color the balls for each prefix? You need additional information:

• the array is circular, so you need to know if person $$$1$$$ has already a colored ball (that wasn't given by person $$$n$$$)
• while processing person $$$i$$$, you need to know if he has already received a colored ball
• you need to know if at least a person in the prefix has given $$$0$$$ balls to the neighbor on his right

Let's do an example of transition. If you want to color one of the balls of person $$$i$$$ and give some balls to his neighbor, including a colored ball, the number of ways can be calculated by fixing the number $$$m$$$ of balls given. For a fixed $$$m$$$, the answer is $$$m(a_i-m)$$$. The total number of ways is $$$\sum_{m=1}^{a_i} m(a_i-m) = a_i \cdot \sum_{m=1}^{a_i} m - \sum_{m=1}^{a_i} m^2 = \frac{a_i^2(a_i+1)}{2} - \frac{a_i(a_i+1)(2a_i+1)}{6}$$$. The other transitions are similar.

Use the trick. Let $$$dp_{i, j, k, l}$$$ the number of ways the people numbered from $$$1$$$ to $$$i$$$ can give and color balls, considering that

• the person $$$1$$$ has colored $$$j$$$ balls ($$$0 \leq j \leq 1$$$)
• the neighbor to the right of person $$$i$$$ has received $$$k$$$ colored balls ($$$0 \leq k \leq 1$$$)
• if $$$l = 0$$$, no person has given $$$0$$$ balls to his neighbor; if $$$l = 1$$$, at least one person has given $$$0$$$ balls to his neighbor.

The transitions can be found by iterating on the answers of the following questions:

• must the person $$$i$$$ color one of his balls?
• how many balls can $$$i$$$ give to his neighbor?
• must the person $$$i$$$ give a colored ball to his neighbor?

An example of transition (with formulas) is in Hint 4.

The answer is the sum of $$$dp_{n, j, 1-j, 1}$$$ (because person $$$1$$$ must have exactly $$$1$$$ colored ball, and there must be a person who has given $$$0$$$ balls to his neighbor).

Complexity: $$$O(n)$$$.

This problem seems very similar to "calculate the number of swaps required to get an array $$$b$$$ (with distinct elements) starting from $$$a$$$", but there can be equal elements.

Some swaps are useless (which ones?)

It's useless to swap equal characters. For each character, can you get its final position in the reversed string? (Treat equal characters as distinguishable values)

Note that it's useless to swap equal characters (in fact, the string doesn't change and you waste $$$1$$$ move).
Consider each character $$$c$$$ from 'a' to 'z' separately. The relative order of a subsequence of equal characters doesn't change in the optimal solution.
Hence, the $$$i$$$-th leftmost occurrence of $$$c$$$ in the initial string should move to the $$$i$$$-th leftmost occurrence of $$$c$$$ in the reversed string. If you number each character of the string, treating equal characters as distinguishable values, you can calculate the array $$$a$$$ such that $$$a_i$$$ is the final position of $$$s_i$$$ in the reversed string, and the result is the number of inversions of $$$a$$$.
For example, consider the string acabcaaababbcb. The reversed string is bcbbabaaacbaca. The positions of 'b' in the reversed string are $$${1, 3, 4, 6, 11}$$$, and this means that $$$a_4 = 1, a_9 = 3, \dots, a_{14} = 11$$$ (in fact, $$$a = [5, 2, 7, 1, 10, 8, 9, 12, 3, 14, 4, 6, 13, 11]$$$).

If there are two adjacent elements, it's optimal to remove them.

In Proof 1, we've seen that a swap changes the number of inversions by $$$1$$$. Here, you should calculate something slightly different than the number of inversions.

Consider two pairs of equal elements. What should their relative positions be at the end?

The answer is $$$n + \text{# of pairs of integers (x, y) such that their order in the array is xyxy}$$$. Can you calculate it efficiently?

Consider two pairs of equal elements (integers $$$(x, y)$$$). If you want to remove those pairs, in some moment their relative positions should be $$$\text{xyyx}$$$ (or $$$\text{yxxy}$$$). Instead, if the relative positions are $$$\text{xyxy}$$$, you want to swap two of these elements. If $$$k$$$ is the number of pairs of integers $$$(x, y)$$$ such that their order in the array is $$$\text{xyxy}$$$, do you need exactly $$$k$$$ swaps (and $$$n$$$ removals)? It turns out that the answer is yes. Sketch of a proof:

• if there are no $$$\text{xyxy}$$$, there are two adjacent equal elements (so you can remove them)
• if you remove two adjacent equal elements, the number of $$$\text{xyxy}$$$ doesn't change
• if there are no adjacent equal elements, there is a $$$\text{xyxy}$$$ such that two of these characters are adjacent (so you can swap them and eliminate a $$$\text{xyxy}$$$)

Now let's calculate $$$k$$$. Let $$$(x_i, y_i)$$$ be $$$n$$$ pairs such that $$$x_i < y_i, a_{x_i} = a_{y_i}$$$: now you have to calculate, for each $$$j$$$, how many indices $$$i$$$ meet $$$x_i < x_j, x_j < y_i < y_j$$$.
You can use a Fenwick tree that contains how many times each value in $$$[1, n]$$$ occurs as a $$$y_i$$$, considering only pairs such that $$$x_i < x_j$$$.
So, for each $$$j$$$, the queries look like

• $$$res := res + \text{range_sum}(x_j + 1, y_j - 1)$$$
• add $$$1$$$ in the position $$$y_j$$$ of the Fenwick tree

When is the answer -1?

The answer is -1 if there are $$$\geq 2$$$ characters with an odd number of occurrences in the string. Otherwise, can you calculate the optimal final position of each character, like in 1430E - String Reversal?

Once again, it's useless to swap equal characters. So, what's the relative position of equal characters? Which characters should be located in symmetrical positions (i.e. $$$i, n - i + 1$$$)?

The $$$i$$$-th leftmost occurrence and the $$$i$$$-th rightmost occurrence of a character $$$c$$$ in the initial string should move to symmetrical positions in the final string. Let's consider such "symmetrical pairs": what's the optimal relative order of two pairs with distinct letters? What about the letter in the center of the string (if its length is odd)?

If there are $$$\geq 2$$$ characters with an odd number of occurrences in the string, you can't make a palindrome.
Since it's useless to swap equal characters, you know which pairs of characters will stay on symmetrical positions in the final string, i.e. the $$$i$$$-th leftmost occurrence and the $$$i$$$-th rightmost occurrence of any character $$$c$$$.
In the string acabcaaababbc, for example, you can make these "symmetrical pairs": positions $$$(1, 10), (2, 13), (3, 8), (4, 12), (6, 7), (9, 11)$$$. The character in position $$$5$$$ should go in the center of the final string (i.e. position $$$7$$$).
The symmetrical pairs have to be nested inside each other, in some order. Given two symmetrical pairs, which of them should be located outside the other one?
It turns out that the pair that contains the leftmost element should be located outside. In fact, if you want to reach $$$\text{xyyx}$$$, the initial configurations with $$$x$$$ on the left ($$$\text{xxyy}$$$, $$$\text{xyxy}$$$, $$$\text{xyyx}$$$) require $$$2, 1, 0$$$ moves respectively, while reaching $$$\text{yxxy}$$$ requires $$$2, 1, 2$$$ moves respectively.
So you can sort the symmetrical pairs by leftmost element and construct the array $$$a$$$ such that $$$a_i$$$ is the final position of $$$s_i$$$ in the palindromic string (in this case, $$$[1, 2, 3, 4, 7, 5, 9, 11, 6, 13, 8, 10, 12]$$$) and the result is the number of inversions of $$$a$$$.

Which balls can you put at the end?

At the end, you have a ball with number $$$n$$$. For example, if it's white, you have to sort the remaining $$$n-1$$$ white balls and $$$n$$$ black balls. How to calculate efficiently the cost required?

Let $$$dp_{i,j}$$$ be the minimum cost required to sort the first $$$i$$$ white balls (with a number from $$$1$$$ to $$$i$$$) and the first $$$j$$$ black balls (with a number from $$$1$$$ to $$$j$$$). You need transitions in $$$O(1)$$$.

Re-read Proof 2 in this tutorial.

It's easy to prove that the leftmost $$$k$$$ balls in the final arrangement are the white balls with a number in $$$[1, i]$$$ and the black balls with a number in $$$[1, j]$$$ for some $$$i, j$$$ such that $$$i+j = k$$$.
For fixed $$$i, j$$$, let $$$dp_{i,j}$$$ be the minimum cost (considering only such balls). If the last ball is white, you can move it immediately to the end (see Proof 2), and the cost is the number of $$$l$$$ in the initial array such that

• $$$l \leq i - 1$$$ if the ball is white, $$$l \leq j$$$ if the ball is black
• $$$\text{pos}(l) > \text{pos}(i)$$$

Then you spend $$$dp_{i-1,j}$$$ to sort the rest of the array. You can calculate the cost similarly if the last ball is black.
Computing such transitions in $$$O(1)$$$ can be done with simple prefix sums (for example, let $$$ps_{i, c, j}$$$ be the number of balls with position $$$\geq i$$$, color $$$c$$$ and value $$$\leq j$$$).
The result is $$$dp_{n,n}$$$.

$$$\text{LCM}$$$ is quite uncomfortable. Rephrase the problem using $$$\text{GCD}$$$.

Use $$$\text{LCM}(a, b) = \frac{ab}{\text{GCD}(a, b)}$$$. Now you can solve the problem separately for each $$$k = \text{GCD}(a_i, a_j)$$$ ($$$i \neq j$$$). What are the optimal values of $$$a_i, a_j$$$?

For a fixed $$$k$$$, it's optimal to choose the minimum $$$a_i, a_j$$$ that are multiples of $$$k$$$. How to find those minimums? Does it matter if, actually, $$$\text{GCD}(a_i, a_j) > k$$$?

You have to minimize $$$\frac{a_ia_j}{\text{GCD}(a_i, a_j)}$$$ ($$$i \neq j$$$).

For each $$$k = \text{GCD}(a_i, a_j)$$$, $$$a_i, a_j$$$ have to be multiples of $$$k$$$. Note that, if $$$\text{GCD}(a_i, a_j) = h > k$$$, you end up calculating a wrong value of $$$\text{LCM}(a_i, a_j)$$$, but it doesn't matter because it's never optimal; you will calculate the correct (smaller) value of the $$$\text{LCM}$$$ while considering $$$\text{GCD}(a_i, a_j) = h$$$ .

So, for each $$$k$$$, it's enough to consider the smallest $$$a_i, a_j$$$ that are multiples of $$$k$$$. You can find them like in Solution 1 or Solution 3 (instead of the sum, you should keep the two minimum elements).

Once again, use $$$\text{GCD}$$$ instead of $$$\text{LCM}$$$ and solve the problem for each fixed $$$\text{GCD}$$$.

Let $$$\displaystyle g(k) = \sum_{k | a_i, k | a_j, i < j}a_ia_j$$$. Can you calculate it with Solution 1 (or Solution 3)?

Let $$$g_1(k) = \sum_{k | a_i}a_i$$$, $$$g_2(k) = \sum_{k | a_i}a_i^2$$$ (they can be calculated with Solution 1 or Solution 3). Then, $$$g(k) = \frac{1}{2}(g_1(k)^2 - g_2(k))$$$ (it's easy to prove).

Now the problem is that some products $$$a_ia_j$$$ considered in $$$g(k)$$$ may have $$$\text{GCD}(a_i, a_j) > k$$$. How to fix that?

You have to calculate the sum of $$$\frac{a_ia_j}{\text{GCD}(a_i, a_j)}$$$ ($$$i < j$$$).

For each $$$k = \text{GCD}(a_i, a_j)$$$, $$$a_i, a_j$$$ have to be multiples of $$$k$$$. Let's start by calculating $$$\displaystyle g(k) = \sum_{k | a_i, k | a_j, i < j}a_ia_j$$$, then we will remove pairs such that $$$\text{GCD}(a_i, a_j) > k$$$.

Let $$$g_1(k) = \sum_{k | a_i}a_i$$$, $$$g_2(k) = \sum_{k | a_i}a_i^2$$$ (they can be calculated with Solution 1 or Solution 3, using $$$b_i = a_i$$$ and $$$b_i = a_i^2$$$ respectively). Note that $$$\displaystyle g_2(k) = \sum_{k | a_i, k | a_j}a_ia_j$$$. Then, $$$g(k) = \frac{1}{2}(g_1(k)^2 - g_2(k))$$$.

Finally, let $$$\displaystyle f(k) = \sum_{\text{GCD}(a_i, a_j) = k, i < j}a_ia_j$$$. Let's calculate $$$f(k)$$$ for $$$k$$$ from $$$m$$$ to $$$1$$$. Then, $$$f(k) = g(k) - \sum_{i=2}^{\lfloor m/k \rfloor} f(ik)$$$ (it's the same idea as the Solution 1).

Suppose you want the last number of the blackboard to be $$$k$$$. Is there an optimal order of operations (such that, if $$$k$$$ can be reached, you can also reach it with that order of operations)?

It's optimal to use all the $$$\text{GCD}$$$ operations before all the $$$\text{MIN}$$$ operations. So, how does the last number of the blackboard look like?

The last number of the blackboard is $$$\leq \min(a_i)$$$, and it's the $$$\text{GCD}$$$ of a subset of $$$a$$$. How to count the possible final numbers?

If $$$k$$$ is the $$$\text{GCD}$$$ of some subset, it's also the $$$\text{GCD}$$$ of $$$S_k = \{a_i \mid k | a_i\}$$$. How to calculate the $$$\text{GCD}$$$ of $$$S_k$$$ for each $$$k$$$?

It's easy to prove that the final number on the blackboard is $$$\leq \min(a_i)$$$, and it's the $$$\text{GCD}$$$ of a subset of $$$a$$$.

For a fixed $$$k$$$, a necessary and sufficient condition for it to be the $$$\text{GCD}$$$ of a subset is to be the $$$\text{GCD}$$$ of the "worst-case" subset, namely $$$S_k = \{a_i \mid k | a_i\}$$$. Let $$$f(k)$$$ be such $$$\text{GCD}$$$.

Initially, all $$$f(k)$$$ are set to $$$0$$$. For each $$$a_i$$$, iterate over its divisors $$$j$$$ (with Solution 2) and set $$$f(j) := \text{GCD}(f(j), a_i)$$$

The answer is the number of $$$k \leq \min(a_i)$$$ such that $$$f(k) = k$$$.

Suppose that you can use $$$x$$$ operations of type $$$1$$$ and $$$y$$$ operations of type $$$2$$$. Try to reorder the operations in such a way that $$$a$$$ becomes the minimum possible.

You should use operations of type $$$2$$$ first, then moves of type $$$1$$$. How many operations do you need in the worst case? ($$$a = 10^9$$$, $$$b = 1$$$)

You need at most $$$30$$$ operations. Iterate over the number of operations of type $$$2$$$.

Notice how it is never better to increase $$$b$$$ after dividing ($$$\lfloor \frac{a}{b+1} \rfloor \le \lfloor \frac{a}{b} \rfloor$$$).

So we can try to increase $$$b$$$ to a certain value and then divide $$$a$$$ by $$$b$$$ until it is $$$0$$$. Being careful as not to do this with $$$b<2$$$, the number of times we divide is going to be $$$O(\log a)$$$. In particular, if you reach $$$b \geq 2$$$ (this requires at most $$$1$$$ move), you need at most $$$\lfloor \log_2(10^9) \rfloor = 29$$$ moves to finish.

Let $$$y$$$ be the number of moves of type $$$2$$$; we can try all values of $$$y$$$ ($$$0 \leq y \leq 30$$$) and, for each $$$y$$$, check how many moves of type $$$1$$$ are necessary.

Complexity: $$$O(\log^2 a)$$$.

If we notice that it is never convenient to increase $$$b$$$ over $$$6$$$, we can also achieve a solution with better complexity.

You can make a $$$k$$$-similar array by assigning $$$a_i = x$$$ for some $$$l \leq i \leq r$$$ and $$$1 \leq x \leq k$$$.

How many $$$k$$$-similar arrays can you make if $$$x$$$ is already equal to some $$$a_i$$$ ($$$l \leq i \leq r$$$)?

How many $$$k$$$-similar arrays can you make if either $$$x < a_l$$$ or $$$x > a_r$$$?

How many $$$k$$$-similar arrays can you make if none of the previous conditions holds?

Let's consider each value $$$x$$$ from $$$1$$$ to $$$k$$$.

• If $$$x < a_l$$$, you can replace $$$a_l$$$ with $$$x$$$ (and you get $$$1$$$ $$$k$$$-similar array). There are $$$a_l-1$$$ such values of $$$x$$$.
• If $$$x > a_r$$$, you can replace $$$a_r$$$ with $$$x$$$ (and you get $$$1$$$ $$$k$$$-similar array). There are $$$k-a_r$$$ such values of $$$x$$$.
• If $$$a_l < x < a_r$$$, and $$$x \neq a_i$$$ for all $$$i$$$ in $$$[l, r]$$$, you can either replace the rightmost $$$b_i$$$ which is less than $$$x$$$, or the leftmost $$$b_i$$$ which is greater than $$$x$$$ (and you get $$$2$$$ $$$k$$$-similar arrays). There are $$$(a_r - a_l + 1) - (r - l + 1)$$$ such values of $$$x$$$.
• If $$$x = a_i$$$ for some $$$i$$$ in $$$[l, r]$$$, no $$$k$$$-similar arrays can be made.

The total count is $$$(a_l-1)+(k-a_r)+2((a_r - a_l + 1) - (r - l + 1))$$$, which simplifies to $$$k + (a_r - a_l + 1) - 2(r - l + 1)$$$.

Complexity: $$$O(n + q)$$$.

Let $$$\lfloor \frac{a}{b} \rfloor = a~\mathrm{mod}~b = k$$$. Is there an upper bound for $$$k$$$?

$$$k \leq \sqrt x$$$. For a fixed $$$k$$$, can you count the number of special pairs such that $$$a \leq x$$$ and $$$b \leq y$$$ in $$$O(1)$$$?

We can notice that, if $$$\lfloor \frac{a}{b} \rfloor = a~\mathrm{mod}~b = k$$$, then $$$a$$$ can be written as $$$kb+k$$$ ($$$b > k$$$). Since $$$b > k$$$, we have that $$$k^2 < kb+k = a \leq x$$$. Hence $$$k \leq \sqrt x$$$.

Now let's count special pairs for any fixed $$$k$$$ ($$$1 \leq k \leq \sqrt x$$$). For each $$$k$$$, you have to count the number of $$$b$$$ such that $$$b > k$$$, $$$1 \leq b \leq y$$$, $$$1 \leq kb+k \leq x$$$. The second condition is equivalent to $$$1 \leq b \leq x/k-1$$$.

Therefore, for any fixed $$$k > 0$$$, the number of special pairs ($$$a\leq x$$$; $$$b \leq y$$$) is $$$max(0, min(y,x/k-1) - k)$$$. The result is the sum of the number of special pairs for each $$$k$$$.

Complexity: $$$O(\sqrt x)$$$.

Brute force doesn't work (even if you optimize it): there are relatively few solutions.

There may be very few possible values of $$$b_{i,j}$$$, if $$$b_{i-1,j}$$$ is fixed. The problem arises when you have to find a value for a cell with, e.g., $$$4$$$ fixed neighbors. Try to find a possible property of the neighbors of $$$(i, j)$$$, such that at least a solution for $$$b_{i,j}$$$ exists.

The least common multiple of all integers from $$$1$$$ to $$$16$$$ is less than $$$10^6$$$.

Build a matrix with a checkerboard pattern: let $$$b_{i, j} = 720720$$$ if $$$i + j$$$ is even, and $$$720720+a_{i, j}^4$$$ otherwise. The difference between two adjacent cells is obviously a fourth power of an integer. We choose $$$720720$$$ because it is $$$\operatorname{lcm}(1, 2, \dots, 16)$$$. This ensures that $$$b_{i, j}$$$ is a multiple of $$$a_{i, j}$$$, because it is either $$$720720$$$ itself or the sum of two multiples of $$$a_{i, j}$$$.

Complexity: $$$O(nm)$$$.

What happens if you can't swap coins?

Let $$$dp_i$$$ be the maximum score that you can reach after $$$dist(1, i)$$$ moves if there is a red coin on node $$$i$$$ after step $$$3$$$. However, after step $$$2$$$, the coin on node $$$i$$$ may be either red or blue. Try to find transitions for both cases.

If you consider only the first case, you solve the problem if there are no swaps. You can greedily check the optimal location for the blue coin: a node $$$j$$$ such that $$$dist(1,i) = dist(1,j)$$$ and $$$a_j$$$ is either minimum or maximum.

Instead, if the coin on node $$$i$$$ is blue after step $$$2$$$, the red coin is on node $$$j$$$ and you have to calculate $$$\max(dp_{parent_j} + |a_j - a_i|)$$$ for each $$$i$$$ with a fixed $$$dist(1, i)$$$. How?

Divide the nodes in groups based on the distance from the root. Then, for each $$$dist(1, i)$$$ in increasing order, calculate $$$dp_i$$$ — the maximum score that you can reach after $$$dist(1, i)$$$ moves if there is a red coin on node $$$i$$$ after step $$$3$$$. You can calculate $$$dp_i$$$ if you know $$$dp_j$$$ for each $$$j$$$ that belongs to the previous group. There are two cases:

• if after step $$$2$$$ the coin on node $$$i$$$ is red, the previous position of the red coin is fixed, and the blue coin should reach either the minimum or the maximum $$$a_j$$$ among the $$$j$$$ that belong to the same group of $$$i$$$;

• if after step $$$2$$$ the coin on node $$$i$$$ is blue, there is a red coin on node $$$j$$$ ($$$dist(1, i) = dist(1, j)$$$), so you have to maximize the score $$$dp_{parent_j} + |a_j - a_i|$$$.
  This can be done efficiently by sorting the $$$a_i$$$ in the current group and calculating the answer separately for $$$a_j \leq a_i$$$ and $$$a_j > a_i$$$; for each $$$i$$$ in the group, the optimal node $$$j$$$ either doesn't change or it's the previous node.
  Alternatively, you can notice that $$$dp_{parent_j} + |a_j - a_i| = \max(dp_{parent_j} + a_j - a_i, dp_{parent_j} + a_i - a_j)$$$, and you can maximize both $$$dp_{parent_j} + a_j - a_i$$$ and $$$dp_{parent_j} + a_i - a_j$$$ greedily (by choosing the maximum $$$dp_{parent_j} + a_j$$$ and $$$dp_{parent_j} - a_j$$$, respectively). In this solution, you don't need to sort the $$$a_i$$$.

The answer is $$$\max(dp_i)$$$.

Complexity: $$$O(n)$$$ or $$$O(n\log n)$$$.

Why isn't the answer $$$2^{n-1}$$$? What would you overcount?

Find a dp with time complexity $$$O(n^2\log n)$$$.

Let $$$dp_{i, j}$$$ be the number of hybrid prefixes of length $$$i$$$ and sum $$$j$$$. The transitions are $$$dp_{i, j} \rightarrow dp_{i+1, j+b_i}$$$ and $$$dp_{i,j} \rightarrow dp_{i+1,b_i}$$$. Can you optimize it to $$$O(n\log n)$$$?

For each $$$i$$$, you can choose either $$$a_i = b_i$$$ or $$$a_i = b_i - \sum_{k=1}^{i-1} a_k$$$. If $$$\sum_{k=1}^{i-1} a_k = 0$$$, the two options coincide and you have to avoid overcounting them.

This leads to an $$$O(n^2\log n)$$$ solution: let $$$dp_i$$$ be a map such that $$$dp_{i, j}$$$ corresponds to the number of ways to create a hybrid prefix $$$[1, i]$$$ with sum $$$j$$$. The transitions are $$$dp_{i, j} \rightarrow dp_{i+1, j+b_i}$$$ (if you choose $$$b_i = a_i$$$, and $$$j \neq 0$$$), $$$dp_{i,j} \rightarrow dp_{i+1,b_i}$$$ (if you choose $$$b_i = \sum_{k=1}^{i} a_k$$$).

Let's try to get rid of the first layer of the dp. It turns out that the operations required are

• move all $$$dp_j$$$ to $$$dp_{j+b_i}$$$
• calculate the sum of all $$$dp_j$$$ in some moment
• change the value of a $$$dp_j$$$

and they can be handled in $$$O(n\log n)$$$ with "Venice technique".

$$$dp$$$ is now a map such that $$$dp_j$$$ corresponds to the number of ways to create a hybrid prefix $$$[1, i]$$$ such that $$$\sum_{k=1}^{i} a_k - b_k = j$$$. Calculate the dp for all prefixes from left to right. if $$$b_i = a_i$$$, you don't need to change any value of the dp; if $$$b_i = \sum_{k=1}^{i} a_k$$$, you have to set $$$dp_{\sum_{k=1}^{i} -b_k}$$$ to the total number of hybrid arrays of length $$$i-1$$$.

Complexity: $$$O(n\log n)$$$.

Solve some small tests on paper.

Suppose that the answer is NO. Which property should hold for the array $$$a$$$?

Solve on paper

3 3
1 1 1
1 1 1
1 1 1

3 3
2 1 1
1 1 2
2 1 2

How does the xor of all the array change after every operation?