Tree Quiz (Problem K, The 2024 ICPC Asia Pacific Championship, 2024)
Time limit: 4 sec
Author: William Gan

# Problem Statement

Your friend wants to quiz you. You are given a *rooted tree* with $n$ nodes, numbered from $1$ to $n$. For every node $i$, its parent is node $p_i$, except for the *root* (the node without a parent) which has $p_i = 0$. Node $u$ is an *ancestor* of node $v$ if either $u = v$, or node $u$ is an ancestor of the parent of node $v$ (if it exists).

We say that node $z$ is a *common ancestor* of nodes $x$ and $y$ if node $z$ is an ancestor of both nodes $x$ and $y$. We say that node $z$ is the *lowest common ancestor* of nodes $x$ and $y$ if it is a common ancestor of nodes $x$ and $y$, and every common ancestor of nodes $x$ and $y$ is also an ancestor of node $z$. We denote the lowest common ancestor of nodes $x$ and $y$ by $\mathrm{LCA}(x, y)$. In particular, $\mathrm{LCA}(x, x) = x$.

Your friend would like to run the following pseudocode:
```
let L be an empty array
for x = 1 to n
    for y = 1 to n
        append ((x - 1) * n * n + (LCA(x, y) - 1) * n + (y - 1)) to L
    sort L in non-decreasing order
```
Your friend has $q$ questions, numbered from $1$ to $q$. In question $j$, you are given an integer $k_j$ and asked to find the $k_j$-th element of the array $L$. Note that $L$ is 1-indexed, so the indices range from $1$ to $n^2$, inclusive. To pass the quiz, you have to answer all of the questions.

# Input

The first line of input contains two integers $n$ and $q$ ($1 \le n \le 100\,000; 1 \le q \le 100\,000$). The second line contains $n$ integers $p_1, p_2, \ldots, p_n$ ($0 \le p_i \le n$ for all $i$). It is guaranteed that the given values represent a rooted tree. Each of the next $q$ lines contains an integer. The $j$-th line contains the string $k_j$ ($1 \le k_j \le n^2$).

# Output

For each question in order, output an integer representing the answer to the question.

# Sample Input #1
```
5 3
3 0 2 2 3
1
```
# Sample Output #1
```
0
82
124
18
25
```
**Explanation for the sample input/output #1**

The tree in the input is illustrated by Figure K.1.

# Figure K.1: Illustration of the tree in sample input #1.

![Figure K.1: Illustration of the tree in sample input #1](icpc-apac-2024-quiz-1.png)

The elements of *L* are  
(0, 6, 8, 12, 14, 30, 31, 32, 33, 34, 56, 58, 60, 62, 64, 80, 81, 82, 84, 93, 106, 108, 110, 112, 124).


Sample inputs:
5 3
3 0 2 2 3
1
18
25

Sample outputs:
0
82
124

# Solution Analysis
Analysis author: Jonathan Irvin Gunawan

The integers in $L$ can be represented as tuples $(x, \text{LCA}(x, y), y)$, where the tuples are ordered by $x$, ties are broken by $\text{LCA}(x, y)$, and further ties are broken by $y$. Since each $i$ has $n$ tuples with $i$ as its first element, the value of $x$ for the $k$-th element of $L$ is simply $\left\lceil \frac{k}{n} \right\rceil$. Next, we want to find the $((k - 1) \mod n + 1)$-th tuple with $x$ as its first element. For simplicity, let $m = (k - 1) \mod n + 1$.

Let $S(i)$ be the set of nodes in the subtree of node $i$, and $A = \{a_1, a_2, \dots\}$ be the ancestors of node $x$ in order from $x$ (i.e., $a_1 = x$ and node $a_{i+1}$ is the parent of node $a_i$). Let $f(a_i) = |S(a_i) \setminus S(a_{i-1})|$ (i.e., the number of nodes in the subtree of node $a_i$ but not in the subtree of node $a_{i-1}$), and $f(z) = 0$ for $z \notin A$. This means $f(z)$ is the number of tuples with $x$ and $z$ as the first and second elements respectively.

```
              a₃ = 2
                /  \
         a₂ = 4     3
           /    \    \        
x = a₁ = 1       5    8
        /  \
        6   7
```
- $f(a_1) = f(1) = 3$ (nodes 1, 6, and 7)
- $f(a_2) = f(4) = 2$ (nodes 4 and 5)
- $f(a_3) = f(2) = 3$ (nodes 2, 3, and 8)
- $f(3) = f(5) = f(6) = f(7) = 0$

The $m$-th tuple that we are looking for has $l$ as $\text{LCA}(x, y)$ where $\sum_{i=1}^{l} f(i) \ge m$ and $\sum_{i=1}^{l-1} f(i) < m$. We can find this value using a segment tree that stores the values of $f(i)$. For each edge $(u, v)$, by changing the value of $x$ from $u$ to $v$, only the values of $f(u)$ and $f(v)$ change. Therefore, we can keep the segment tree for each $x$ using persistent segment tree in $O(n \log n)$ preprocessing time.

Let $m' = m - \sum_{i=1}^{l-1} f(i)$. We want to find the $m'$-th tuple with $x$ and $l$ as its first and second elements respectively. Let $p$ be the index of $A$ such that $a_p = l$. Thus, we want to find the $m'$-th node (in increasing order) in $S(a_p) \setminus S(a_{p-1})$.

We can do this by using another segment tree. Let $G_i$ be a segment tree that stores $G_i(j) = 1$ if node $i$ is an ancestor of node $j$, or $G_i(j) = 0$ otherwise. To compute $G_i$, we can take $G_h$ as the base, where node $h$ is the child of node $i$ with maximum $|S(h)|$, and change the values of $G_i(j)$ for all nodes $j \in (S(i) \setminus S(j))$ individually. With another persistent segment tree, we can keep all segment trees $G$ in $O(n \log^2 n)$ preprocessing time.

With $G_{a_p}$ and $G_{a_{p-1}}$ which represents the subtree of node $a_p$ and node $a_{p-1}$ respectively, we can find the answer $y$ that we are looking for. It is the unique integer that satisfies the following two inequalities:

$$
\left( \sum_{i=1}^{y} G_{a_p}(i) \right) - \left( \sum_{i=1}^{y} G_{a_{p-1}}(i) \right) \ge m'
$$

$$
\left( \sum_{i=1}^{y-1} G_{a_p}(i) \right) - \left( \sum_{i=1}^{y-1} G_{a_{p-1}}(i) \right) < m'
$$

This solution runs in $O(n \log^2 n) + q \log^2(n))$ time for preprocessing the persistent segment trees and answering all questions.


Special judge (C++ scorer):
/**
 * Default validator from https://github.com/Kattis/problemtools
 * Modified (as small as possible) to tcframe format by jonathanirvings.
 */

#include <fstream>
#include <iostream>
#include <string>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cassert>
#include <cmath>
#include <cstdarg>

std::ifstream judgein, judgeans, conans;
FILE *judgemessage = NULL;
FILE *diffpos = NULL;
int judgeans_pos, con_pos;
int judgeans_line, con_line;

void ac() {
  puts("AC");
  exit(0); 
}

void wa() {
  puts("WA");
  exit(0);
}

void wrong_answer(const char *err, ...) {
  wa();
}

void judge_error(const char *err, ...) {
  va_list pvar;
  va_start(pvar, err);
  // If judgemessage hasn't been set up yet, write error to stderr
  if (!judgemessage) judgemessage = stderr;
  vfprintf(judgemessage, err, pvar);
  fprintf(judgemessage, "\n");
  assert(!"Judge Error");
}

bool isfloat(const char *s, double &val) {
  char trash[20];
  double v;
  if (sscanf(s, "%lf%10s", &v, trash) != 1) return false;
  val = v;
  return true;
}

template <typename Stream>
void openfile(Stream &stream, const char *file, const char *whoami) {
  stream.open(file);
  if (stream.fail()) {
    judge_error("%s: failed to open %s\n", whoami, file);
  }
}

FILE *openfeedback(const char *feedbackdir, const char *feedback, const char *whoami) {
  std::string path = std::string(feedbackdir) + "/" + std::string(feedback);
  FILE *res = fopen(path.c_str(), "w");
  if (!res) {
    judge_error("%s: failed to open %s for writing", whoami, path.c_str());
  }
  return res;
}

const char *USAGE = "Usage: %s judge_in judge_ans contestant_ans";

int main(int argc, char **argv) {
  if(argc < 4) {
    judge_error(USAGE, argv[0]);
  }
  openfile(judgein, argv[1], argv[0]);
  openfile(judgeans, argv[2], argv[0]);
  openfile(conans, argv[3], argv[0]);

  bool case_sensitive = false;
  bool space_change_sensitive = false;
  bool use_floats = false;
  double float_abs_tol = -1;
  double float_rel_tol = -1;

  for (int a = 4; a < argc; ++a) {
    if        (!strcmp(argv[a], "case_sensitive")) {
      case_sensitive = true;
    } else if (!strcmp(argv[a], "space_change_sensitive")) {
      space_change_sensitive = true;
    } else if (!strcmp(argv[a], "float_absolute_tolerance")) {
      if (a+1 == argc || !isfloat(argv[a+1], float_abs_tol))
        judge_error(USAGE, argv[0]);
      ++a;
    } else if (!strcmp(argv[a], "float_relative_tolerance")) {
      if (a+1 == argc || !isfloat(argv[a+1], float_rel_tol))
        judge_error(USAGE, argv[0]);
      ++a;
    } else if (!strcmp(argv[a], "float_tolerance")) {
      if (a+1 == argc || !isfloat(argv[a+1], float_rel_tol))
        judge_error(USAGE, argv[0]);
      float_abs_tol = float_rel_tol;
      ++a;
    } else {
      judge_error(USAGE, argv[0]);
    }
  }
  use_floats = float_abs_tol >= 0 || float_rel_tol >= 0;

  judgeans_pos = con_pos;
  judgeans_line = con_line = 1;
   
  std::string judge, team;
  char trash[20];
  while (true) {
    // Space!  Can't live with it, can't live without it...
    while (isspace(judgeans.peek())) {
      char c = (char)judgeans.get();
      if (space_change_sensitive) {
        int d = conans.get();
        if (c != d) {
          wrong_answer("Space change error: got %d expected %d", d, c);
        }
        if (d == '\n') ++con_line;
        ++con_pos;
      }
      if (c == '\n') ++judgeans_line;
      ++judgeans_pos;
    }
    while (isspace(conans.peek())) {
      char d = (char)conans.get();
      if (space_change_sensitive) {
        wrong_answer("Space change error: judge out of space, got %d from team", d);
      }
      if (d == '\n') ++con_line;
      ++con_pos;
    }

    if (!(judgeans >> judge))
      break;

    if (!(conans >> team)) {
      wrong_answer("User EOF while judge had more output\n(Next judge token: %s)", judge.c_str());
    }
     
    double jval, tval;
    if (use_floats && isfloat(judge.c_str(), jval)) {
      if (!isfloat(team.c_str(), tval)) {
        wrong_answer("Expected float, got: %s", team.c_str());
      }
      if(!(fabs(jval - tval) <= float_abs_tol) && 
         !(fabs(jval - tval) <= float_rel_tol*fabs(jval))) {
        wrong_answer("Too large difference.\n Judge: %s\n Team: %s\n Difference: %le\n (abs tol %le rel tol %le)", 
               judge.c_str(), team.c_str(), jval-tval, float_abs_tol, float_rel_tol);
      }
    } else if (case_sensitive) {
      if (strcmp(judge.c_str(), team.c_str()) != 0) {
        wrong_answer("String tokens mismatch\nJudge: \"%s\"\nTeam: \"%s\"", judge.c_str(), team.c_str());
      }
    } else {
      if(strcasecmp(judge.c_str(), team.c_str()) != 0) {
        wrong_answer("String tokens mismatch\nJudge: \"%s\"\nTeam: \"%s\"", judge.c_str(), team.c_str());
      }
    }
    judgeans_pos += judge.length();
    con_pos += team.length();
  }

  if (conans >> team) {
    wrong_answer("Trailing output:\n%s", team.c_str());
  }

  ac();
}