PROBLEM LINK:

Practice
Contest

Author: admin
Tester: Kevin Atienza
Editorialist: Ajay K. Verma
Russian Translator: Sergey Kulik
Mandarian Translator: Gedi Zheng

DIFFICULTY:

Medium

PREREQUISITES:

Combinatorics, Modular Arithmetic

PROBLEM:

Given two integers $N$ and $K$, find how many triplets of strings over $K$ alphabets exist, such that length of all strings is bounded by $N$, and no string in a triplet is a prefix of another string of the triplet.

EXPLANATION:

We compute the number of triplets of strings of bounded length, and then subtract the ones in which one string is a prefix of another. For the sake of brevity, we define a new notation $\prec$ to represent the "prefix" relationship, i.e., "$X$ is a prefix of $Y$" can be represented by $X \prec Y$.

Total Number of Triplets:

There are exactly $K^i$ strings of length $i$. Hence, the number of strings with length not exceeding $N$ will be:
$A = K + K^2 + K^3 + \cdots + K^N$
$A = K (K^N - 1 / (K - 1))$

If the modular inverse of $(K - 1)$ with respect to the given prime $p = 10^9 + 7$ is $x$, then the value of $A \bmod p$ would be $(K (K^N - 1) x) \bmod p$.

Both the modular inverse, and power of a number can be compute in $O (\log N)$ time. Hence, we can compute the value of $A$ in $O (\log N)$ time.

There are $A$ strings of length not exceeding $N$. Therefore, the number of triplets would be $A^3$. Next, we count how many of these triplets are bad triplets, i.e., the ones in which one string is a prefix of another.

Triplets of the Form $X \prec Y \prec Z$:

In this section we count the number of triplets which have strings forming a chain under the relationship $\prec$, i.e., if the triplet is $(P, Q, R)$, then one of the following is true:
$P \prec Q \prec R$
$P \prec R \prec Q$ 
$Q \prec P \prec R$
$Q \prec R \prec P$ 
$R \prec P \prec Q$
$R \prec Q \prec P$

Note that for some triplets more than one of these conditions might be true, e.g., $(X, X, X)$ will satisfy all $6$ conditions. Hence, we consider four categories, so that each triplet is counted exactly once.

1) $X = Y = Z$:
There will be $A$ triplets of this form, as we can choose the string $X$ in $A$ ways, and the other two will be equal to it. These triplets will satisfy all $6$ conditions, i.e., each triplet will be repeated $6$ times.

2) $X \prec Y = Z$ and $X \neq Y$:
For a given string $Y$ of length $r$, there are exactly $(r - 1)$ possible values of string $X$, the ones which are proper substring of $Y$. Hence, the number of such triplets would be:
$K^2 + 2 K^3 + 3 K^4 + \cdots + (N - 1) K^N$.

This is the sum of an arithmetic-geometric series, and can be represented explicitly as shown here. The computation of this sum also requires the computation of power and modular inverse, and hence can be done in $O (\log N)$ time.

These triplets will satisfy two of the above conditions ($X \prec Y \prec Z$, and $X \prec Z \prec Y$).

3) $X = Y \prec Z$, and $Y \neq Z$:
Number of these triplets will be the same as the one computed in (2). This is because once we pick a string $Z$ of length $r$, then $Y$ has exactly $(r - 1)$ choices.

These triplets also satisfy two of the above conditions ($X \prec Y \prec Z$, and $Y \prec X \prec Z$).

4) $X \prec Y \prec Z$, $X \neq Y$, and $Y \neq Z$:
If we pick a string $Z$ of length $r$, then we only need to choose the lengths of $X$ and $Y$. The lengths of $X$ and $Y$ must be distinct and lie between $1$ and $(r - 1)$. Hence, the number of $(X, Y)$ pairs would be the number of ways of choosing two distinct numbers between $1$ and $(r - 1)$, which is given by $(r - 1) \choose 2$. Therefore the number of such triplets would be:
$K^3 + 3 K^4 + 6 K^5 + \cdots + {{N - 1} \choose 2} K^N$

Again this is the sum of arithmetic-geometric sequence of second order, and can be computed using the similar methods. The computation will take $O (\log N)$ time.

These triplets will satisfy exactly one of the above conditions.

Hence, the number of chain triplets would be $6$ ((number of category 1 triplets)/$6$ + (number of category 2 triplets)/$2$ + (number of category 3 triplets)/$2$ + (number of category 4 triplets)/$1$).

Triplets of the Form "$X \prec Y$, $Z$ does not share prefix relationship with $Y$":

In this section we count the number of triplets in which exactly two strings satisfy the "prefix" relationship, i.e., if the triplet is $(P, Q, R)$, then one of the following is true:
$P \prec Q$, $R$ has no prefix relation with $Q$,
$Q \prec P$, $R$ has no prefix relation with $P$,
$P \prec R$, $Q$ has no prefix relation with $R$,
$R \prec P$, $Q$ has no prefix relation with $P$,
$Q \prec R$, $P$ has no prefix relation with $R$,
$R \prec Q$, $P$ has no prefix relation with $Q$,

Once again it is possible that one triplet satisfy more than one condition, however, we should count it only once.

1) $X = Y$, $Z$ has no prefix relation with $Y$: 
If we have a string $Y$ of length $r$, then there are exactly $r$ prefix strings of $Y$. If a string $W$ has $Y$ as a proper prefix, then $W$ can be written as $W = Y + U$, where $U$ is a non-empty string of length at most $(N - r)$. Hence, the number of such string $W$ will be $K + K^2 + \cdots + K^{N - r}$. Since $Z$ is neither a prefix of $Y$, nor has $Y$ as its prefix, the number of possible values of $Z$ would be:
$(K + K^2 + \cdots + K^N) - (K + K^2 + \cdots + K^{N - r}) - r$
$= (K^{N - r + 1} + K^{N - r + 2} + \cdots + K^N - r)$.

Also number of ways of choosing a $r$ length string $Y$ is $K^r$. Hence, the number of such triplets $(X, Y, Z)$, with $Y$ being a $r$-length string will be:
$K^r (K^{N - r + 1} + K^{N - r + 2} + \cdots + K^N - r)$
$= K^{N + 1} * (1 + K + K^2 + \cdots + K^{r - 1}) - r * K^r$

If we take the sum of this expression over all possible values of $r$, we get the number of triplets $(X, Y, Z)$ satisfying the above criteria, and it will be:
$K^{N + 1} (N + (N - 1) K + (N - 2) K^2 + \cdots + K^{N - 1}) - (K + 2 K^2 + 3 K^3 + \cdots + N K^N)$

This is also a sum of arithmetic-geometric sequence, and can be computed in logarithmic time. These triplets satisfy two of the above conditions ($X \prec Y$, with $Z$ having no relationship, and $Y \prec X$, with $Z$ having no relationship).

2) $X \prec Y$, $X \neq Y$, and $Z$ has no prefix relation with $Y$:
Once again if we pick a string $Y$ of length $r$, then $X$ will have $(r - 1)$ choices, and Z will have $(K^{N - r + 1} + K^{N - r + 2} + \cdots + K^N - r)$ choices. Hence, the number of such triplets with $Y$ being of length $r$ would be:

$r K^r (K^{N - r + 1} + K^{N - r + 2} + \cdots + K^N - r)$
$= r K^{N + 1} (1 + K + K^2 + \cdots + K^{r - 1}) - r^2 K^r $

If we sum it over all possible values of $r$, we would get the number of triplets $(X, Y, Z)$ satisfying the above condition, which would be (after some simplifications):

${(N + 1) \choose 2} K^{N + 1} (1 + K + K^2 + \cdots + K^N) - K^{N + 1} ({2 \choose 2} K + {3 \choose 2} K^2 + \cdots + {N \choose 2} K^{N - 1})$
$ - (1^2 K + 2^2 K^2 + 3^2 K^3 + \cdots + N^2 K^N)$.

The number of triplets where exactly two strings satisfy the "prefix" relationship will be $6$ ( (number of triplets of first category)/$2$ + (number of triplets of second category)/$1$).

Now that we have computed the number of bad triplets, we can subtract them from the total number of triplets to get the number of good ones.

Time Complexity:

$O (\log N)$

AUTHOR'S AND TESTER'S SOLUTIONS:

Author's solution will be put up soon.
Tester's solution will be put up soon.