Problem: double_stars (2023 round3)

Patrick's big sister Sam is visiting for Halloween, and Sandy wants to impress the two Stars by dressing as an astronaut. Unfortunately her regular outfit already resembles an astronaut, so she'll have to go as a computer scientist instead. Sandy knows a lot about [double stars](https://en.wikipedia.org/wiki/Double_star) in astronomy, but will now have to learn about double stars in graph theory.

There is an unrooted tree with \(N\) nodes, and an edge between nodes \(i\) and \(P_i\) (for each \(i \ge 2\)). A *double star* is a tree centered at some edge \(u \leftrightarrow v\), where \(u\) and \(v\) each have \(x\) chains connected to them, each chain consisting of \(y\) nodes (\(x, y \ge 1\)). Alternatively, you can think of a double star as two identical star graphs (each with \(x\) chains of length \(y\)) connected by their centers. Thus, a distinct double star is specified by a tuple \((u, v, x, y)\), where \(u \lt v\). 

Please help Sandy count the number of distinct tuples \((u, v, x, y)\) for which the tree contains the double star described by that tuple as a subgraph.

# Constraints

\(1 \leq T \leq 110\)
\(2 \leq N \leq 1{,}000{,}000\)
\(1 \leq P_i \leq N\)

The sum of \(N\) across all test cases is at most \(8{,}000{,}000\).

# Input Format

Input begins with an integer \(T\), the number of test cases. Each case has two lines. The first contains the integer \(N\), and the second contains the \(N-1\) integers \(P_2\) through \(P_N\).

# Output Format

For the \(i\)th test case, print "`Case #i:` " followed by the number of distinct double star tuples \((u, v, x, y)\), where \(u \lt v\).

# Sample Explanation
The first sample tree is depicted below. There are \(5\) distinct double star tuples: \((1, 2, 1, 1)\), \((1, 2, 1, 2)\), \((1, 2, 2, 1)\), \((1, 4, 1, 1)\) and \((2, 6, 1, 1)\).

{{PHOTO_ID:1573931960080150|WIDTH:350}}

The second sample tree is depicted below. There are \(5\) distinct double star tuples: \((1, 2, 1, 1)\), \((1, 2, 1, 2)\), \((1, 8, 1, 1)\), \((1, 8, 2, 1)\) and \((2, 7, 1, 1)\).

{{PHOTO_ID:1286569645378322|WIDTH:500}}

In the third sample tree, no double stars exist because arm lengths of \(0\) are not allowed \((y \ge 1 )\).


Sample input:
3
8
1 1 1 2 2 6 4
10
1 8 8 1 1 2 1 7 7
2
1


Sample output:
Case #1: 5
Case #2: 5
Case #3: 0


Solution: For each vertex \(u\), and for each incident edge \(e\), we first have to calculate the longest distance to another vertex coming from this edge \(e\). It can be observed that for each vertex, there can be at most \(\mathcal{O}(\sqrt N)\) distinct distances.

We'll build a table \(d(u, v)\) storing the longest arm length starting with \(u \to v \to ...\), which can be computed in \(\mathcal{O}(N)\) using two DFS's. For each vertex \(v\), we have a list of pairs \((d, x)\) indicating that from vertex \(v\) there are \(x\) different edges incident to \(v\) in which the furthest distance you can reach is \(d\). We'll store this in a table \(f(u)\) holding a map of \((d, x)\), denoting that \(d\) is the greatest distance of an arm \(u \to v\) for \(x\) edges \(v\). \(f(u)\) can be computed using just the table \(d(u, v)\),

From here, the problem can be solved using the two-pointer technique. We'll want consider every edge \(u \leftrightarrow v\), and efficiently generate every \((x, y)\) that defines a double star tree with \(u \leftrightarrow v\) as its center. For each \((u, v, x, y)\), we iterate through \(f(u)\) and \(f(v)\) with two-pointers, with \(d'\) initially set to \(1\):
  - For the first \((d, x)\), take the sum of \(\min(\text{degree}(u) - 1, \text{degree}(v) - 1)\) over every \([d', d]\) (on those \(d\)'s, we can make a double star with that number of arms).
  - Remove the first and decrease \(x\) from \(\text{degree}(u)\) or \(\text{degree}(v)\), depending on what the pair belongs to.
  - Set \(d' = d + 1\).

The are \(\mathcal{O}(N)\) possible edges for the center of the star, each of which we'll examine \(\mathcal{O}(\sqrt N)\) distinct distances using the two-pointers. The overall time complexity is \(\mathcal{O}(N \sqrt N)\).


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

// Solution:
// - For each edge u-v, find d(u, v) = longest arm length starting with u-v-..
// - For each edge u-v, consider all double stars (x, y) with u-v as the center.
// - Iterate through f[u] and f[v] with 2-pointer, with d' initially 1:
//   - For the first (d, x), sum min(degree(u) - 1, degree(v) - 1) for every
//     [d', d] (those d's, we can make 2-star with that number of arms).
//   - Remove the first and decrease x from degree(u) or degree(v), depending
//     on what the pair belongs to.
//   - Set d' = d + 1

vector<vector<int>> adj;
vector<vector<int>> d;  // d[u][i] = greatest distance reachable from adj[u][i]
vector<map<int, int>> f;  // f[u] = map of (d, x), representing d is the
                          // greatest distance of an arm (u, v) for x edges v

int pre(int u, int par = -1, int l = 0) {
  int mai = 0;
  for (int i = 0; i < (int)adj[u].size(); ++i) {
    int v = adj[u][i];
    d[u].push_back(-1);
    if (v == par) {
      continue;
    }
    int ret = pre(v, u);
    mai = max(mai, ret);
    d[u][i] = ret;
  }
  return 1 + mai;
}

void pre2(int u, int par = -1, int farthest_up = 0) {
  vector<int> pref(adj[u].size()), suf(adj[u].size());
  for (int i = 0; i < (int)adj[u].size(); ++i) {
    pref[i] = max(i ? pref[i - 1] : 0, d[u][i]);
  }
  for (int i = (int)adj[u].size() - 1; i >= 0; --i) {
    suf[i] = max(i + 1 < (int)adj[u].size() ? suf[i + 1] : 0, d[u][i]);
  }
  for (int i = 0; i < (int)adj[u].size(); ++i) {
    int v = adj[u][i];
    if (v == par) {
      d[u][i] = farthest_up;
      continue;
    }
    // Farthest going down here.
    int tmp = 1 + max(i ? pref[i - 1] : 0,
                      i + 1 < (int)adj[u].size() ? suf[i + 1] : 0);
    pre2(v, u, max(1 + farthest_up, tmp));
  }
}

long long rec(int u, int par = -1) {
  long long ans = 0;
  for (int i = 0; i < (int)adj[u].size(); ++i) {
    int v = adj[u][i];
    if (v == par) {
      // Calculate for DST centered at (u, par).
      f[u][d[u][i]]--;

      auto it1 = f[v].begin();
      auto it2 = f[u].begin();
      int tot1 = (int)adj[v].size() - 1;
      int tot2 = (int)adj[u].size() - 1;

      int l = 1;
      while (it1 != f[v].end() || it2 != f[u].end()) {
        if (it1 != f[v].end() &&
            (it2 == f[u].end() || it1->first < it2->first)) {
          ans += (long long)(it1->first - l + 1) * min(tot1, tot2);
          tot1 -= it1->second;
          l = it1->first + 1;
          it1++;
        } else {
          ans += (long long)(it2->first - l + 1) * min(tot1, tot2);
          tot2 -= it2->second;
          l = it2->first + 1;
          it2++;
        }
      }
      // assert(tot1 == 0 && tot2 == 0);
      f[u][d[u][i]]++;
      continue;
    }
    f[u][d[u][i]]--;
    ans += rec(v, u);
    f[u][d[u][i]]++;
  }
  return ans;
}

long long solve() {
  int N;
  cin >> N;
  adj.assign(N + 1, {});
  d.assign(N + 1, {});
  f.assign(N + 1, {});
  for (int i = 0, p; i < N - 1; i++) {
    cin >> p;
    adj[i + 2].push_back(p);
    adj[p].push_back(i + 2);
  }
  pre(1);
  pre2(1);
  for (int u = 1; u <= N; ++u) {
    for (auto dd : d[u]) {
      f[u][dd]++;
    }
  }
  return rec(1);
}

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