PROBLEM LINK:
Author: Vasia Antoniuk
Tester: Mahbubul Hasan and Sunny Aggarwal
Editorialist: Balajiganapathi Senthilnathan
Russian Translator: Sergey Kulik
Mandarian Translator: Minako Kojima
DIFFICULTY:
Medium
PREREQUISITES:
Z function, combinatorics
PROBLEM:
Given a string and a query k, output the number of ways we can choose k equal substrings from the string.
SHORT EXPLANATION
Let cnt[i] be the number of different substrings that occur in the string exactly i times. Then the answer for a particular k ( \le n) will be :
\sum_{i = k}^{n} cnt[i] * \binom{i}{k} we can calculate cnt[i] using suffix array/Z function.
EXPLANATION:
Let cnt[i] be the number of different substrings that occur in the string exactly i times.
For example, for the string “ababa”, the array will be:
cnt[1] = 4 -> {“ababa”, “abab”, “baba”, “bab”} occur exactly once
cnt[2] = 4 -> {“aba”, “ab”, “ba”, “b”} occur exactly twice
cnt[3] = 1 -> {“a”} occur exactly thrice
cnt[4] = 0
cnt[5] = 0
Now suppose we want the answer for k = 2, i.e. number of ways to choose 2 equal strings. How do we calculate this using the cnt array?
cnt[1] is of no use as we need 2 equal string.
cnt[2] -> There are 4 different substrings that occur 2 times. So, we can add 4 to the answer.
cnt[3] -> There is only 1 substring that occurs 3 times. But we need only 2 times, so we can choose any 2 of the 3. So, \binom{3}{2} = 3
Summing them, for k = 2 we get 7.
Now let us focus on calculating the cnt array.
To calculate this let us loop through all suffixes of S.
For the suffix P = S[i..N], let us calculate the array Z[i..N] where Z[i] is the maximal equal substring starting from i that matches a prefix of S[i..N].
For example, for the above string consider i = 1, the whole string is the suffix
“ababa”.
Here the Z array would be
Z[1] = 5 -> maximum prefix matching of “ababa” and “ababa”
Z[2] = 0 -> maximum prefix matching of “ababa” and “baba”
Z[3] = 3 -> maximum prefix matching of “ababa” and “aba”
Z[4] = 0 -> maximum prefix matching of “ababa” and “ba”
Z[5] = 1 -> maximum prefix matching of “ababa” and “a”
How is this array useful for calculating cnt?
From the above array we can conclude that we have 1 substring which can be choosen atleast 3 times. How? Because there are 3 entries in the above array that are greater than or equal to 1.
Similarly we can deduce for 2 (2 times), 3 (2 times), 4 (1 time ) and 5(1 time).
So we increment cnt[3] once, cnt[2] twice and cnt[1] twice.
We do this for all suffix of S. The only thing left is to observe that cnt[i] now actually has how many different strings can be choosen atleast i times. We want it to be exactly i times. This is simple: we just subtract from cnt[i] (cnt[i + 1] + cnt[i + 2] … + cnt[N]).
We can compute the Z function in O(n) time. See here for how to do that.
Time Complexity:
There are N suffixes and we take O(N) time for each suffix so total time is:
$O(N^2).