Solving AGC 058 D with multivariate generating functions

Revision en1, by adamant, 2022-08-15 00:04:50

Hi everyone!

Today, I participate in AtCoder contest, which I quite rarely do. Of course, my performance was quite poor, but I stumbled with quite nice idea along the way. To those who didn't read it, the problem goes as follows:

AGC 058 — Yet Another ABC String. How many ternary strings are there such that they contain exactly $$$A$$$ characters a, exactly $$$B$$$ characters b and exactly $$$C$$$ characters c, and do not have any sub-string that is a cyclic shift of abc?

My approach is very different from the one in the editorial, and seems insightful to me, so I decided to share it in detail.

Besides, I can't reject the request by Um_nik to explain it further :D

Walks on finite automata

We may construct a directed graph, such that each edge is marked by a, b or c, and so that the string is valid if and only if there is a walk from the starting vertex so that concatenation of characters on traversed edges form the string. The graph looks like this:


Graph of allowed walks

In the graph with have $$$7$$$ vertices:

  • Starting vertex in the middle;
  • Vertices $$$a$$$, $$$b$$$ and $$$c$$$ denote the last character and the fact that pre-last character is not important;
  • Vertices $$$ab$$$, $$$bc$$$ and $$$ca$$$ denote the last two characters.

As you see, going from $$$ab$$$ by $$$c$$$, from $$$bc$$$ by $$$a$$$ or from $$$ca$$$ by $$$b$$$ is not allowed. The graph may be further simplified to only $$$3$$$ vertices:


Reduced graph of allowed walks

Indeed, we can't traverse from vertices $$$ab$$$, $$$bc$$$ and $$$ca$$$ in one another, so we may as well remove them completely.

Multivariate generating function

Let $$$F(a, b, c)$$$ be the generating function on the number of allowed walks, so that $$$[a^A b^B c^C]F(a, b, c)$$$ is the number of allowed walks that have $$$A$$$ characters a, $$$B$$$ characters b and $$$C$$$ characters c. Then it may be expressed as

$$$ F = a F_a + b F_b + c F_c, $$$

where $$$F_a$$$, $$$F_b$$$ and $$$F_c$$$ are the generating functions for walks starting in $$$a$$$, $$$b$$$ and $$$c$$$ correspondingly.

From the graph structure it follows that

$$$\begin{cases} F_a &=& (a + ba) F_a &+& b^2 F_b &+& c F_c &+& 1, \\ F_b &=& (b + cb) F_b &+& c^2 F_c &+& a F_a &+& 1, \\ F_c &=& (c + ac) F_c &+& a^2 F_a &+& b F_b &+& 1. \end{cases}$$$

This may be expressed as a system of linear equations, and the whole answer is given as

$$$ F = \begin{pmatrix} a & b & c \end{pmatrix} \begin{pmatrix} 1 - a - ba & -b^2 & -c \\ -a & 1-b-cb & -c^2 \\ -a^2 & -b & 1-c-ac \end{pmatrix}^{-1} \begin{pmatrix} 1 \\ 1 \\ 1 \end{pmatrix}. $$$

I'll be honest here, I don't exactly know the way to compute it manually, but matrixcalc suggests that the answer is

$$$ \large F(a, b, c) = \frac{ab+ac+bc-a-b-c}{a+b+c-2abc-1}. $$$

Note that the computation is not exactly pretty:


Matrixcalc "suggestion"

What to do with generating function?

Well, first of all multiplying both parts by the denominator we get that

$$$ F(a, b, c) = (a+b+c) F(a, b, c) - 2abc F(a, b, c) + F_0(a, b, c), $$$

where $$$F_0(a, b, c) = ab + ac + bc - (a+b+c)$$$ is insignificant to higher terms of $$$F$$$.

From the formula above it follows that

$$$ F[A][B][C] = F[A-1][B][C] + F[A][B-1][C] + F[A][B][C-1] - 2F[A-1][B-1][C-1], $$$

where $$$F[A][B][C] = [a^A b^B c^C]F(a, b, c)$$$ is the answer to the problem. Not exactly the result that is obvious from the statement, is it?

However, the problem requires the solution in $$$O(A+B+C)$$$ and, luckily, it is also possible to find it from the generating function!

Let $$$G(a,b,c) = \frac{1}{a+b+c-2abc-1}$$$. Note that

$$$\begin{align} [a^Ab^B c^C] F =& [a^{A-1}b^{B-1}c^C] G + [a^{A-1}b^{B}c^{C-1}] G + [a^{A}b^{B-1}c^{C-1}] G \\ &-[a^{A-1}b^{B}c^C] G - [a^{A}b^{B-1}c^{C}] G - [a^{A}b^{B}c^{C-1}] G, \end{align}$$$

hence the computation of $$$[a^A b^B c^C] F$$$ can be reduced to $$$6$$$ computations of similar thing in $$$G$$$.

And to compute $$$[a^A b^B c^C] G$$$ you should note that

$$$ \large G = \sum\limits_{t=0}^\infty (a+b+c-2abc)^t = \sum\limits_{i,j,k,l} \binom{i+j+k+l}{i,j,k,l} a^i b^j c^k (-2abc)^l. $$$

Among all the terms on the right, only $$$\min(A, B, C)$$$ contribute to $$$[a^A b^B c^C] G$$$ and we may find them all by iterating over $$$l$$$.

Here, $$$\binom{i+j+k+l}{i,j,k,l} = \frac{(i+j+k+l)!}{i!j!k!l!}$$$ is a multinomial coefficient.

Big thanks to Golovanov399 for sharing this approach!

Tags tutorial, generating function, dfa

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en9 English adamant 2022-08-15 04:17:40 7
en8 English adamant 2022-08-15 00:36:44 348
en7 English adamant 2022-08-15 00:34:39 1
en6 English adamant 2022-08-15 00:27:21 91
en5 English adamant 2022-08-15 00:24:36 86
en4 English adamant 2022-08-15 00:23:57 179
en3 English adamant 2022-08-15 00:21:26 18
en2 English adamant 2022-08-15 00:07:01 102
en1 English adamant 2022-08-15 00:04:50 5041 Initial revision (published)