Problem: third_trie (2022 round3)

You have \(N\) [tries](https://en.wikipedia.org/wiki/Trie). The \(i\)th trie has \(M_i\) nodes and \(M_i - 1\) edges, with node \(1\) being the root, and node \(j\)'s parent being \(P_{i,j}\), for \(j = 2..M_i\). Edge \(P_{i,j} \to j\) is labeled with a lowercase letter \(C_{i,j}\).

We say a string is *in* a trie if it can be formed by concatenating the letters along some path starting at the root and ending at some node (either internal or leaf).

For every possible unordered triplet of these \(N\) tries, you would like to know: how many strings exist which are in at least one of the three tries?

Please print the sum of answers over all \(N*(N-1)*(N-2)/6\) possible queries.


# Constraints

\(1 \le T \le 100\)
\(3 \le N \le 100\)
\(1 \le M_i \le 1{,}000{,}000\)
\(1 \le P_{i,j} \lt j\)
\(C_{i,j} \in \{\)`'a'`\(,\, \ldots, \) `'z'`\(\}\)

The sum of \(N\) across all cases is at most \(2{,}000\).
No node will have two direct outgoing edges with the same edge label.
There are at most \(1{,}000{,}000\) total trie nodes in a given test case.
The sum of nodes across all test cases is at most \(13{,}000{,}000\).


# Input Format

Input begins with a single integer \(T\), the number of test cases. For each test case, there is first a line containing a single integer \(N\), the number of tries. Then, \(N\) descriptions of tries follow. For the \(i\)th trie, there is first a line containing a single integer \(M_i\), followed by \(M_i - 1\) more lines, the \(j\)th of which contains the integer \(P_{i,j+1}\), followed by a space, followed by the letter \(C_{i,j+1}\).


# Output Format

For the \(i\)th test case, output a single line containing `"Case #i: "` followed by a single integer, the sum of the query answers across all triplets of tries.


# Sample Explanation

In the first case, there are \(N = 3\) tries, depicted below:

{{PHOTO_ID:502671685043623|WIDTH:650}}

There is only one possible unordered triplet of tries to consider, for which there are \(5\) strings in at least one trie of the triplet: "", "`a`", "`aa`", "`b`", and "`c`".

In the second case, there are \(N = 4\) tries, depicted below:

{{PHOTO_ID:1118217749060271|WIDTH:650}}

There are four triplets of tries to consider, three of which include the last trie and one of which doesn't. In the three that does, there are \(4\) strings in at least one trie: "", "`a`", "`aa`", and "`aaa`". In the one that doesn't, the are \(2\) strings: "" and "`a`". The final answer is \(3*4 + 1*2 = 14\).

In the third case, there are \(N = 4\) tries, depicted below:

{{PHOTO_ID:1007664223967563|WIDTH:650}}

The four triplets and strings contributing to the answer are:
* \((\)trie \(1\), trie \(2\), trie \(3)\): "", "`a`", "`b`", "`c`"
* \((\)trie \(1\), trie \(2\), trie \(4)\): "", "`a`", "`b`", "`ab`", "`abc`"
* \((\)trie \(1\), trie \(3\), trie \(4)\): "", "`a`", "`c`", "`ab`", "`abc`"
* \((\)trie \(2\), trie \(3\), trie \(4)\): "", "`a`", "`b`", "`c`", "`ab`", "`abc`"

The final answer is \(4 + 5 + 5 + 6 = 20\).


Sample input:
3
3
3
1 a
1 b
3
1 a
2 a
2
1 c
4
2
1 a
2
1 a
2
1 a
4
1 a
2 a
3 a
4
2
1 a
2
1 b
2
1 c
4
1 a
2 b
3 c


Sample output:
Case #1: 5
Case #2: 14
Case #3: 20


Solution: Each trie node represents a distinct string \(s\), so the number of times that \(s\) will be present in a triplet of tries is independent of all other strings. This leads to the observation that the answer can be calculated independently over each string, as long as we know the number of occurrences of \(s\) across all tries.

Supposing a string \(s\) is in \(k\) tries, we like to know it's contribution to the answer: for how many out of all \(\binom{N}{3}\) possible triples of tries is it in at least one of the tries? The answer is \(\binom{N}{3}\) minus the number the number of triples that doesn't have this node in any of the three tries, equal to \(\binom{N - k}{3}\).

All that remains is computing each pair of \((s, k_s\)). We can use DFS across all \(N\) tries to build a "mega" trie consisting of all \(N\) tries overlayed on top of each other, with the count at each node of how many times the corresponding string appears. This running time of the DFS will be proportional to the total number of nodes. Finally, we can DFS again to walk the entire megatrie, adding the contribution of \(\binom{N}{3} - \binom{N - k}{3}\) to the answer for each \(k\).

Alternatively, we can hash all strings as we walk through all the tries, maintaining a table of the counts for each hash, calculating the answer at the end just the same.


Code: #include <iostream>
#include <vector>
using namespace std;

const int ALPHA_SIZE = 26;

struct node {
  int count;
  vector<node *> child;

  node(int _count = 0) : count(_count), child(ALPHA_SIZE, nullptr) {}

  void merge(node *n) {
    count++;
    for (int i = 0; i < ALPHA_SIZE; i++) {
      if (n->child[i] != nullptr) {
        if (child[i] == nullptr) {
          child[i] = new node();
        }
        child[i]->merge(n->child[i]);
      }
    }
  }
};

class trie {
  node *root;

  template <typename Func>
  static void walk(node *n, Func f) {
    f(n->count);
    for (int i = 0; i < ALPHA_SIZE; i++) {
      if (n->child[i] != nullptr) {
        walk(n->child[i], f);
      }
    }
  }

  static void clean_up(node *n) {
    for (int i = 0; i < ALPHA_SIZE; i++) {
      if (n->child[i] != nullptr) {
        clean_up(n->child[i]);
      }
    }
    delete n;
  }

 public:
  trie() : root(new node()) {}
  ~trie() { clean_up(root); }
  void merge(node *n) { root->merge(n); }

  template <typename Func>
  void walk(Func f) const {
    walk(root, f);
  }
};

long long solve() {
  int N, p;
  char c;
  cin >> N;
  vector<vector<node>> tries(N);
  for (int i = 0, M; i < N; i++) {
    cin >> M;
    auto &T = tries[i];
    T.resize(M);
    for (int j = 1; j < M; j++) {
      cin >> p >> c;
      T[--p].child[c - 'a'] = &T[j];
    }
  }
  trie merged;
  for (auto &T : tries) {
    merged.merge(&T[0]);
  }
  auto n_choose_3 = [](long long n) {
    return n < 3 ? 0 : n * (n - 1) * (n - 2) / 6;
  };
  long long ans = 0;
  merged.walk([&](int k) {
    // This node is in k tries. For how many out of all N choose 3 possible
    // triples (T1, T2, T3) of tries is this node in at least one of the tries?
    // The answer is N choose 3 minus the number of triples that doesn't have
    // this node in any of the three tries, equal to (N - k) choose 3.
    ans += n_choose_3(N) - n_choose_3(N - k);
  });
  return ans;
}

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