問題

$n$ 頂点からなる木構造のグラフと、整数 $1, \cdots , n$ が1個ずつランダムに入った順列 $p$ が与えられる。

ここで、グラフ内の頂点 $u, v$ の距離を $\mathrm{dist}(u, v)$ とする。

$q$ 組の整数 $l, r, x$ が与えられるので、そのそれぞれについて $\mathrm{dist}(1, v) = \mathrm{dist}(1, x) + \mathrm{dist}(x, v)$ が成立するような $v$ が ${p}_{l}, {p}_{l + 1}, \cdots , {p}_{r}$ の中に含まれるか判定せよ。

入力

まず最初の1行目に、テストケースの個数を表す整数 $t$ が与えられる。

その後、 $t$ 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

最初の1行目には、整数 $n, q$ が与えられる。
次の $n - 1$ 行には、グラフの $i$ 番目の辺が接続されている2頂点を表す整数 ${u}_{i}, {v}_{i}$ が1行ずつ与えられる。
次の1行には、順列 ${p}_{i}$ の値が順に与えられる。
次の $q$ 行には、それぞれ $l, r, x$ の値が1行ずつ与えられる。

条件

$1 \leq t \leq {10}^{4}$
- すべてのテストケースにおける $n, q$ の値の合計はいずれも ${10}^{5}$ を超えない。
- 与えられる辺の情報について、グラフは木構造であることが保証される。
$1 \leq l \leq r \leq n$
$1 \leq x \leq n$

出力

それぞれのテストケースについて、 $q$ 行に渡り、それぞれの $(l, r, x)$ に対して条件を満たすものが存在するならば YES を、存在しないならば NO を出力すること。

解法

まず、頂点 $x$ に対して $\mathrm{dist}(1, v) = \mathrm{dist}(1, x) + \mathrm{dist}(x, v)$ が成立する頂点 $v$ とはどのようなものなのか、ということを考えます。

与えられるグラフは木構造であるため、頂点 $1$ を根として考えます。このとき例えば、以下の図のようなグラフが与えられたとします。

このグラフに対して、例えば頂点 $x, v$ が以下のような位置にあるとすると、条件式が成立します。

一方で、以下の2つの図のような位置関係の場合は、条件式は成立しません。

これらの図から、「頂点 $1$ から頂点 $v$ へと辿り着く最短経路の中に、頂点 $x$ が含まれること」が与えられた式の成立条件であると言えます。これを頂点 $x$ 目線で言い換えると、「頂点 $x$ から奥にある頂点（頂点 $x$ も含む）の中に、頂点 $v$ があること」ということになります。

今回の問題では、「与えられた頂点 $x$ に対して、条件を満たす頂点 $v$ が頂点 ${p}_{l}, {p}_{l + 1}, \cdots , {p}_{r}$ の中に含まれるのか」を調べなければなりません。そこで、「頂点 $x$ よりも奥にある頂点 $v$ に対して、 ${p}_{i} = v$ であるすべての $i$ 」をすべて列挙するということを考えます。

以下の図のように「頂点 $1$ から見て端にある頂点から、順に辺を辿る」という方法で、 $n$ 頂点分を効率的に巡ることができます。

従って、C++での set などを用いることで、すべての頂点 $x$ に対して、条件を満たす $v$ の ${p}_{i} = {v}$ となる $i$ の値をすべて保存することができます。

このときに、頂点 $x$ に対して保存された値の中に、 ${l}, \cdots , {r}$ のものが含まれているかどうかを調べる方法について考えます。

先程の値の保存を set で行ったとすると、 lower_bound *1 を用いて考えます。

lower_bound は「指定した値 $x$ に対して $\leq x$ を満たす最初のイテレータを返す」関数なので、 $x = l$ としてまずはイテレータをもらいます。その返ってきたイテレータに対して、「それが .end() ではない」かつ「その要素が $r$ 以下である」ということを満たせば、条件を満たす値が $l, \cdots, r$ の中に存在する、ということが言えます。

以上がこの問題の解法となりますが、このままでは set の中身を作成していく際に、頂点 $1$ のものを作るには $n$ 個分の中身を入れる必要があります。そのため、すべての頂点に対して中身を入れる作業を行うことを考えると、最大で $\displaystyle \frac{1}{2}n(n-1)$ 個分の中身を入れる必要が出てきてしまい、 $n \leq {10}^{5}$ より、これは実行時間制限とメモリ制限を超えてしまう可能性があります。

頂点 $1$ に近い頂点であればあるほど、 set に新たに追加しなければならない値が増えます。

先程の探索の図より、ある頂点について一度探索したら、それ以降その頂点についての探索は行いません。

従って、現在見ている頂点よりも深い位置にある頂点での set のすべての要素を、現在見ている頂点での set に移す際に、「現在見ている頂点の set の要素数よりも深い位置にある頂点の set の要素数が多い場合は、「2つの set を入れ替えて、要素数の少ない方から多い方へと値を移していく」という方法を用いることを考えます。

このことにより、現在頂点 $x$ を見ているとすると、頂点 $x$ についての条件を満たす $i$ の値の保存が効率化できると言えます。

ただこの方法では、set の入れ替えを行っているため、頂点 $x$ よりも深い位置にある頂点の set の中身は現実と異なるものとなります。

そこで、この問題ではグラフとともに $l, r, x$ の値の組み合わせも予め与えられるので、「頂点 $x$ についての探索を終えたら、その頂点についてのすべてのクエリの回答を用意しておく」ということを行います。これを行うことにより、ある頂点 $x$ よりも深い位置にある頂点についてのクエリはすべて答えられている状態になるので、それらの set の中身が実際とは異なっていても特に問題は起きないと言えます。

この方法で、 set に値を投入する回数が削減できるので、実行時間制限・メモリ制限を守って解答を出力できると考えられます。

ソースコード

まず、ある頂点 $x$ に対して条件を満たす ${p}_{i} = v$ であるすべての $i$ を保存し、 $x$ に対してのクエリの答えを保存する calc_index() を以下のように実装しました。なお、 n, p, q, graph, queries にはすでに値が入力されているものとして記載しています。

また、すべての $i$ （ただし、 $1 \leq i \leq n$ ）に対して、予め indexes[p[i]] に i が入っているものとします。（すべての頂点 $x$ に対して、 $\mathrm{dist}(1, x) = \mathrm{dist}(1, x) + \mathrm{dist}(x, x)$ が必ず成立するため）

int n, q, p[int(1e5 + 5)];
vector<int> graph[int(1e5 + 5)]; // graph[i] は、頂点 i に接続されているすべての頂点を保存する

bool ans[int(1e5 + 5)]; // ans[i] は i 問目の答えが true なのか false なのかを表す

// クエリ用の struct を作成する
struct query {
  int l, r, i; // l, r は与えられた値で、 i は「何問目のクエリなのか」という情報を表す
};
vector<query> queries[int(1e5 + 5)]; // queries[x] には、頂点 x についての query の内容が入っている

set<int> indexes[int(1e5 + 5)]; // indexes[x] は、頂点 x の条件を満たすすべての p_i = v  である i が入る

// now_node は現在の頂点、 before_node はその1つ前（頂点1に近い側）の頂点を表す
void calc_index(int now_node, int before_node = -1) {
  // 頂点 now_node に接続されたすべての頂点について調査する
  for (int next_node : graph[now_node]) {
    if (next_node == before_node) {
      continue; // before_node については indexes に入れる必要がないので、無視する
    }
    calc_index(next_node, now_node); // 次の頂点について探索する

    // next_node についての探索後、indexes[next_node] に入っている値をすべて indexes[now_node] に移す

    if (indexes[now_node].size() < indexes[next_node].size()) {
      swap(indexes[now_node], indexes[next_node]); // 値が入る側の set の要素数を大きくしておく
    }

    for (int i : indexes[next_node]) {
      indexes[now_node].insert(i); // 要素数が小さい set のすべての要素を、要素数が大きい set に移す
    }
  }

  // x = now_node となっているクエリ全てに回答する
  for (query now_query : queries[now_node]) {
    auto iter = indexes[now_node].lower_bound(now_query.l); // l 以上の最初の要素へのイテレータを受け取る
    if (iter != indexes[now_node].end() && (*iter) <= now_query.r) {
      ans[now_query.i] = true; // 条件を満たすならば true を入れる
    } else {
      ans[now_query.i] = false; // 満たさないならば false を入れる
    }
  }

  return;
}

この calc_index() がある状態で、問題の解答を出力する solve() 関数を以下のように実装しました。なお、コード中の prepare_graph() 関数は、 $n - 1$ 個の辺を構成する頂点 $u, v$ の値を受け取り、 graph にそれを反映するための単純なものです。

int n, q, p[int(1e5 + 5)];
vector<int> graph[int(1e5 + 5)]; // graph[i] は、頂点 i に接続されているすべての頂点を保存する

bool ans[int(1e5 + 5)]; // ans[i] は i 問目の答えが true なのか false なのかを表す

// クエリ用の struct を作成する
struct query {
  int l, r, i; // l, r は与えられた値で、 i は「何問目のクエリなのか」という情報を表す
};
vector<query> queries[int(1e5 + 5)]; // queries[x] には、頂点 x についての query の内容が入っている

set<int> indexes[int(1e5 + 5)]; // indexes[x] は、頂点 x の条件を満たすすべての p_i = v  である i が入る

void solve() {
  // 値の入力
  cin >> n >> q;
  prepare_graph();
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
    indexes[p[i]].insert(i); // 予め indexes[p[i]] に i を入れておく
  }

  // クエリの入力
  for (int i = 1; i <= q; i++) {
    int l, r, x;
    cin >> l >> r >> x;
    queries[x].push_back({l, r, i}); // 頂点 x についてのクエリが i 問目なので、 l, r, i を入れる
  }

  // 頂点 1 が根にあたるので、頂点 1 について探索するようにすれば、すべての頂点が再帰的に探索される
  calc_index(1);

  // 答えの出力
  for (int i = 1; i <= q; i++) {
    if (ans[i]) {
      YES();
    } else {
      NO();
    }
  }
  return;
}