Problem: bohemian_rap-sody (2023 round1)

*Is this a real verse? Is this just fantasy?
Caught in a mixtape, no escape from this melody.
Open your ears, listen to the beats and see.
I'm just a rapper, I need no sympathy,
Because it's easy come, easy go, little high, little low,
Any way the bars flow doesn't really matter to me, to me.*

As an up-and-coming rapper, you're trying to distinguish yourself by writing a song with the most bohemian rhyme scheme!

Your vocabulary consists of \(N\) words, \(W_1\), \(…\), \(W_N\) made up of only lowercase letters. We consider two words to *\(K\)-rhyme* with each other if they have the same suffix of length \(K\). For example, “mixtape” 3-rhymes with “escape”. If two words \(K\)-rhyme for \(K > 1\), they also \((K-1)\)-rhyme.

You are faced with \(Q\) bohemian rap-queries, the \(i\)th of which specifies \((A_i, B_i, K_i)\). You’d like to assign each word \(W_{A_i}\), \(…\), \(W_{B_i}\) to at most one group such that:

* the number of words in each group is unique and nonzero,
* groups only contain words with length at least \(K_i\),
* within each group, all pairs of words \(K_i\)-rhyme with each other, and
* no two words from different groups \(K_i\)-rhyme.

The answer to each query is the maximum possible number of nonempty groups. For each test case, output the sum of the answers over all queries. 

**Constraints**

\(1 \le T \le 20\)
\(1 \le N \le 600{,}000\)
\(\sum |W_i| \le 3{,}000{,}000\)
\(1 \le Q \le 400{,}000\)
\(1 \le A_i \le B_i \le N\)
\(1 \le K_i \le 1{,}000{,}000\)
All words consist of only letters \(\{\texttt{`a'}, ..., \texttt{`z'}\}\)
All words within a test case are distinct.

The sum of \(N\) across all test cases is at most \(2{,}000{,}000\).
The sum of \(Q\) across all test cases is at most \(1{,}000{,}000\).
The sum of word lengths across all test cases is at most \(10{,}000{,}000\).

**Input Format**

Input begins with an integer \(T\), the number of test cases. For each case, there is first a line containing a single integer \(N\). \(N\) lines follow, the \(i\)th of which contains the string \(W_i\). Then, there is a line containing a single integer \(Q\). \(Q\) lines follow, the \(i\)th of which contains \(3\) space-separated integers \(A_i\), \(B_i\), and \(K_i\).

**Output Format**

For the \(i\)th test case, print "`Case #i:` " followed by a single integer, the sum of answers over all queries.

**Sample Explanation**

For the first query in the first test case, we are allowed to choose from all \(W_1, ..., W_8\) words to form \(3\) groups where words of the same group \(1\)-rhyme with each other. One possible choice of the groups is: \(\{\)"low"\(\}\), \(\{\)"sympathy", "fantasy", "melody"\(\}\), \(\{\)"mixtape", "escape"\(\}\) with "come" and "high" going unused. Therefore the answer to the first query is \(3\). The answers to the remaining queries are \(1\), \(2\) and \(2\), with the overall sum of answers for this test case equal to \(8\).

In the second test case, the answers to the \(3\) queries are \([2, 1, 0]\) with the sum being \(3\).



Sample input:
2
8
sympathy
fantasy
melody
mixtape
escape
come
high
low
4
1 8 1
1 8 7
2 5 1
1 5 3
3
a
aa
b
3
1 3 1
1 3 2
1 3 3


Sample output:
Case #1: 8
Case #2: 3


Solution: The problem is equivalent to finding the maximum \(x\) where you can make \(x\) groups with distinct sizes \(1, ..., x\), and all words that belong to a certain group pairwise \(K\)-rhyme with each other.

The first step is to simplify things by reversing all strings so that we're dealing with prefixes of length \(K\) instead of suffixes. Then, order queries by increasing \(K\). For equal \(K\), we sort queries based on [Mo’s algorithm](https://cp-algorithms.com/data_structures/sqrt_decomposition.html#mos-algorithm).

Note that there are \(\sum |W_i|\) possible prefixes across all the words. We can uniquely identify each prefix by a rolling hashing on each word, or by inserting each word into a trie and using the trie node ID or memory address. These are the values for which we are maintaining the frequency within Mo’s algorithm.

At each query being handled by Mo's algorithm, we should be careful to only process active words for the current size \(K\). The overall time complexity is \(\mathcal{O}((\sum |W_i|)*\sqrt N)\). In practice, our model solution solves the full test set in under \(10\) seconds.


Code: #include <algorithm>
#include <cmath>
#include <iostream>
#include <set>
#include <tuple>
#include <vector>
using namespace std;

using int64 = long long;

int64 N, Q, sqrtsz;
vector<int64> freq, ff;

void inc(int64 x) {
  freq[x]++;
  ff[freq[x]]++;
}

void dec(int64 x) {
  ff[freq[x]]--;
  freq[x]--;
}

int64 get_ans() {
  int64 ans = ff[1];
  for (int i = 2; i <= ans; i++) {
    int64 need = ans - (i - 1);
    if (ff[i] < need) {
      ans -= need - ff[i];
    }
  }
  return ans;
}

int64 solve() {
  cin >> N;

  vector<string> vs(N);
  for (int i = 0; i < N; i++) {
    cin >> vs[i];
  }
  int64 sumc = 0;
  int maxt = 0;
  for (int i = 0; i < N; i++) {
    reverse(vs[i].begin(), vs[i].end());
    sumc += vs[i].size();
    maxt = max(maxt, (int)vs[i].size());
  }

  ff = vector<int64>(N + 10);
  freq = vector<int64>(sumc + 10);

  int idt = 0;
  vector<vector<int>> prox(sumc + 10, vector<int>(26));

  for (auto &s : vs) {
    int est = 0;
    for (auto &c : s) {
      int ic = c - 'a';
      if (!prox[est][ic]) {
        idt++;
        prox[est][ic] = idt;
      }
      est = prox[est][ic];
    }
  }

  vector<int64> vest(N);
  set<int> active;
  for (int i = 0; i < N; i++) {
    active.insert(i);
  }
  vector<vector<tuple<int64, int64, int64>>> vq(maxt + 1);

  cin >> Q;
  vector<int64> ans(Q);
  for (int i = 0; i < Q; i++) {
    int64 l, r, k;
    cin >> l >> r >> k;
    l--, r--;
    if (k <= maxt) {
      vq[k].emplace_back(l, r, i);
    }
  }

  // Reset ff only by the number of active! (in reality you only need sqrt).
  // No need to reset freq, because they are all different states!
  for (int k = 1; k <= maxt; k++) {
    for (int i = 0; i <= (int)active.size(); i++) {
      ff[i] = 0;
    }

    for (auto it = active.begin(); it != active.end();) {
      int i = *it;
      if (k > (int)vs[i].size()) {
        it = active.erase(it);
      } else {
        int ic = vs[i][k - 1] - 'a';
        vest[i] = prox[vest[i]][ic];
        it++;
      }
    }
    vector<int> vid(active.begin(), active.end());

    vector<tuple<int64, int64, int64>> vqq;
    vqq.reserve(vq[k].size());
    for (auto &[lc, rc, ia] : vq[k]) {
      lc = lower_bound(vid.begin(), vid.end(), lc) - vid.begin();
      rc = upper_bound(vid.begin(), vid.end(), rc) - vid.begin();
      rc--;
      if (lc <= rc) {
        vqq.emplace_back(lc, rc, ia);
      }
    }
    vq[k] = vqq;
    sqrtsz = sqrt(vid.size()) + 1;
    sort(vq[k].begin(), vq[k].end(),
      [](tuple<int64, int64, int64> &a, tuple<int64, int64, int64> &b) {
        int64 la, ra, ia;
        tie(la, ra, ia) = a;
        int64 lb, rb, ib;
        tie(lb, rb, ib) = b;
        if (la / sqrtsz != lb / sqrtsz) {
          return la < lb;
        }
        return ra < rb;
      });

    int l = 0, r = -1;
    for (auto &[lc, rc, ia] : vq[k]) {
      while (r < rc) {
        r++;
        inc(vest[vid[r]]);
      }
      while (l > lc) {
        l--;
        inc(vest[vid[l]]);
      }
      while (l < lc) {
        dec(vest[vid[l]]);
        l++;
      }
      while (r > rc) {
        dec(vest[vid[r]]);
        r--;
      }
      ans[ia] = get_ans();
    }
  }
  int64 res = 0;
  for (int i = 0; i < Q; i++) {
    res += ans[i];
  }
  return res;
}

int main() {
  int T;
  cin >> T;
  for (int t = 1; t <= T; t++) {
    cout << "Case #" << t << ": " << solve() << endl;
  }
  return 0;
}