Editorial for COCI '21 Contest 1 #4 Set

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

For the first subtask, it is sufficient to try out every triplet of cards and check whether it forms a set in $\mathcal O(n^3 k)$ ~\mathcal O(n^3 k)~.

For the second subtask, one should notice that if we choose two cards, the third card needed for a set is uniquely determined. Namely, for each position, if the corresponding characters in the first two cards are the same, then the third character has to be equal to them, and if they are different, the third character will be the remaining one. With time complexity $\mathcal O(n^2 k)$ ~\mathcal O(n^2 k)~, we can now try out every pair of cards and check whether there exists a third card which will form a set with them. Of course, we should first record in an array of size $3^k$ ~3^k~ for each type of card if it shows up in the input.

From now on, we will assume that the characters in question are $0, 1, 2$ ~0, 1, 2~ instead of $1, 2, 3$ ~1, 2, 3~, so that we can view each card as a number in base $3$ ~3~. Let's try to come up with a simple rule that associates a pair of characters with the third character needed for a set. In other words, we are looking for a rule that acts in the following way:

$\displaystyle (0, 0) \mapsto 0 \quad (1, 1) \mapsto 1 \quad (2, 2) \mapsto 2 \quad (0, 1) \mapsto 2 \quad (0, 2) \mapsto 1 \quad (1, 2) \mapsto 0$

We can notice that $(a, b) \mapsto -(a+b) \bmod 3$ ~(a, b) \mapsto -(a+b) \bmod 3~, or equivalently, the characters $a, b, c$ ~a, b, c~ form a set if and only if $(a+b+c) \bmod 3 = 0$ ~(a+b+c) \bmod 3 = 0~. Now let's make an analogy with the bitwise xor operation. This operation is denoted by $\oplus$ ~\oplus~ and represents addition modulo $2$ ~2~ with the digits in base $2$ ~2~. In this problem, we will consider addition modulo $3$ ~3~ with the digits in base $3$ ~3~, which we will denote by $*$ ~*~. Having in mind the things mentioned above, three cards $a, b, c$ ~a, b, c~ form a set if and only if $a*b*c = 0$ ~a*b*c = 0~, where the cards are thought of as numbers in base $3$ ~3~.

Let's fix some card $c$ ~c~ and try to figure out how many pairs of cards $a$ ~a~ and $b$ ~b~ exist so that $a*b = -c$ ~a*b = -c~. For each individual $c$ ~c~, we can calculate this in $\mathcal O(nk)$ ~\mathcal O(nk)~ so the total complexity is $\mathcal O(n^2 k)$ ~\mathcal O(n^2 k)~, but we will show a way to find the answer for all cards $c$ ~c~ simultaneously in the complexity $\mathcal O(3^k k)$ ~\mathcal O(3^k k)~. If instead of the operation $*$ ~*~ we had the operation $+$ ~+~, the problem could be solved by calculating the desired $+$ ~+~-convolution with fast multiplication of polynomials using FFT. In this problem, we will therefore try to modify this idea to calculate the $*$ ~*~-convolution.

The operations $\oplus$ ~\oplus~ and $*$ ~*~ are very similar and the $*$ ~*~-convolution is calculated in the same manner as the xor-convolution. Thus, what follows is a description of the modification of the 'fast Walsh–Hadamard transformation' (FWHT) to work modulo $3$ ~3~. More about this can be found on this Codeforces blog. (P.S. the day before the contest, another great blog on the topic appeared on Codeforces: link)

At a high level, the idea is the following:

The given deck of cards is represented by a polynomial.
Each term in the polynomial represents a certain type of card.
The coefficients of the polynomial represent the number of times a card appears in the deck. In the beginning, all of the coefficients are either $0$ ~0~ or $1$ ~1~.
We will square the polynomial to get new coefficients which represent the result of the desired convolution. (For now, this corresponds to the $+$ ~+~ operation).
Before multiplying, we will convert the polynomial from coefficient form to point value form as a sequence of calculated values $(x_i, P(x_i))$ ~(x_i, P(x_i))~, which is more desirable for multiplication.
The result of the multiplication should be converted back to coefficient form.

Regular multiplication of polynomials corresponds to the operation $+$ ~+~, that is $(x^a, x^b) \mapsto x^{a+b}$ ~(x^a, x^b) \mapsto x^{a+b}~. We would like to make a modification so that $(x^a, x^b) \mapsto x^{a*b}$ ~(x^a, x^b) \mapsto x^{a*b}~. We need to make two modifications:

Addition should be done separately for each digit.
Addition should be done modulo $3$ ~3~.

Problem 1 can be solved by introducing a polynomial with $k$ ~k~ variables $x_0, x_1, \dots, x_{k-1}$ ~x_0, x_1, \dots, x_{k-1}~. For example, a pair of cards $1123$ ~1123~ and $2321$ ~2321~ (that is $0012$ ~0012~ and $1210$ ~1210~) is represented by a polynomial

$\displaystyle P(x_3, x_2, x_1, x_0) = x_3^0 x_2^0 x_1^1 x_0^2 + x_3^1 x_2^2 x_1^1 x_0^0$

Multiplication now corresponds to addition of the digits separately. Looking at each of the variables separately, the polynomial is of degree at most $2$ ~2~, but when squaring, the degree might grow larger.

To convert a polynomial (which remember has $3^k$ ~3^k~ coefficients) to point-value form, we will calculate the value of the polynomial at $3^k$ ~3^k~ different points. We will choose three values $v_0, v_1, v_2$ ~v_0, v_1, v_2~ and calculate $P(x_{k-1}, \dots, x_0)$ ~P(x_{k-1}, \dots, x_0)~ for each possible combination where $x_i \in \{v_0, v_1, v_2\}$ ~x_i \in \{v_0, v_1, v_2\}~, of which there are also $3^k$ ~3^k~. In the implementation, we will therefore have a transformation $F$ ~F~ that converts a sequence of $3^k$ ~3^k~ coefficients to a sequence of $3^k$ ~3^k~ calculated values.

Since the product of two polynomials has more coefficients than the original polynomials, if we wanted to calculate the true product, we would have had to extend the polynomials to bigger powers, making the new coefficients $0$ ~0~. However, to solve problem 2, that is precisely what we will not do. When doing the inverse transformation $F^{-1}$ ~F^{-1}~ which returns the coefficient form, we will purposefully demand that the result has $3^k$ ~3^k~ coefficients. Additionally, for the values $v_0, v_1, v_2$ ~v_0, v_1, v_2~, we will choose the third roots of unity (both real and complex), that is the numbers $1$ ~1~, $\frac{-1+\sqrt{3}i}{2}$ ~\frac{-1+\sqrt{3}i}{2}~ and $\frac{-1-\sqrt{3}i}{2}$ ~\frac{-1-\sqrt{3}i}{2}~, so that $v_0^3 = v_1^3 = v_2^3 = 1$ ~v_0^3 = v_1^3 = v_2^3 = 1~. The effect of these two things is that the coefficients of larger powers in the resulting product $(x^3, x^4, \dots)$ ~(x^3, x^4, \dots)~ will get added to the coefficients of the smaller powers - precisely with the smallest power that is the same modulo $3$ ~3~ (because of the choice of $v_i$ ~v_i~). Thus, the powers will reduce modulo $3$ ~3~ and we will get exactly the desired $3^k$ ~3^k~ coefficients of the $*$ ~*~-convolution.

Let us illustrate this on an example with $k = 2$ ~k = 2~. For the polynomial, we take

$\displaystyle \begin{align*} P(x_1, x_0) &= a_{00} x_1^0 x_0^0 + a_{01} x_1^0 x_0^1 + a_{02} x_1^0 x_0^2 \\ &+ a_{10} x_1^1 x_0^0 + a_{11} x_1^1 x_0^1 + a_{12} x_1^1 x_0^2 \\ &+ a_{20} x_1^2 x_0^0 + a_{21} x_1^2 x_0^1 + a_{22} x_1^2 x_0^2 \end{align*}$

which can also be written as

$\displaystyle P(x_1, x_0) = Q_0(x_0) \cdot x_1^0 + Q_1(x_0) \cdot x_1^1 + Q_2(x_0) \cdot x_1^2$

where

$\displaystyle Q_i(x_0) = a_{i0} x_0^0 + a_{i1} x_0^1 + a_{i2} x_0^2$

First, we will apply the transformation $F$ ~F~ on each of the polynomials $Q_i$ ~Q_i~ separately. Using the label $w = -\frac 1 2 + \frac{\sqrt 3}{2} i$ ~w = -\frac 1 2 + \frac{\sqrt 3}{2} i~, we have:

$\displaystyle (a_{i0}, a_{i1}, a_{i2}) \mapsto (a_{i0} + a_{i1} + a_{i2}, \quad a_{i0} + w a_{i1} + w^2 a_{i2}, \quad a_{i0} + w^2 a_{i1} + w a_{i2})$

In this way, the sequence of coefficients

$\displaystyle (a_{00}, a_{01}, a_{02}, \quad a_{10}, a_{11}, a_{12}, \quad a_{20}, a_{21}, a_{22})$

turns into a new list of coefficients, which we will label with

$\displaystyle (a'_{00}, a'_{01}, a'_{02}, \quad a'_{10}, a'_{11}, a'_{12}, \quad a'_{20}, a'_{21}, a'_{22})$

We would like to obtain a list that has the following values in order

$\displaystyle P(1, 1), P(1, w), P(1, w^2), \quad P(w, 1), P(w, w), P(w, w^2), \quad P(w^2, 1), P(w^2, w), P(w^2, w^2)$

What remains is to replace the triplets $(a'_{0i}, a'_{1i}, a'_{2i})$ ~(a'_{0i}, a'_{1i}, a'_{2i})~ with new values, in the same manner as we did when calculating the values for $Q_i$ ~Q_i~.

For bigger values of $k$ ~k~, the process is analogous. In each of the $k$ ~k~ iterations, we gradually transform the coefficients in the described way, each time jumping by a larger power of $3$ ~3~.

The inverse transformation $F^{-1}$ ~F^{-1}~ is almost identical to the original, having the formula

$\displaystyle (a, b, c) \mapsto \frac 1 3 (a + b + c, \quad a + w^2 b + w c, \quad a + w b + w^2 c)$

It should be noted that in the implementation, there is no need to make calculations with complex numbers, whose real and imaginary parts are stored with a floating point type. Instead, we can notice that at each moment every number will be of the form $a+bw$ ~a+bw~, where $a$ ~a~ and $b$ ~b~ are whole numbers which fit in long long int. When calculating, it is useful to keep in mind that $w^3 = 1$ ~w^3 = 1~ and $w^2 = -1-w$ ~w^2 = -1-w~.

Comments

There are no comments at the moment.