問題

$n$ 頂点の木があり、頂点は $1, 2, \cdots, n$ と、辺は $1, 2, \cdots, n - 1$ と番号をつけられている。また、辺 $i$ （ただし、 $1 \leq i \leq n - 1$ ）は頂点 ${u}_{i}$ と頂点 ${v}_{i}$ を結んでいる。

$1 \leq l \leq r \leq n - 1$ を満たすすべての整数組 $(l, r)$ に対して、以下の値の総和を求めよ。

頂点 $1$ から出発し、番号が $l$ 以上 $r$ 以下である辺のみを $0$ 本以上辿って到達できる頂点の個数。

入力

まず最初の1行目には、整数 $n$ が与えられる。

次の $n - 1$ 行のうち $i$ 行目には ${u}_{i}, {v}_{i}$ が順に与えられる。

条件

$2 \leq n \leq 2 \times {10}^{5}$
$1 \leq {u}_{i} \lt {v}_{i} \leq n$ （ただし、 $1 \leq i \leq n - 1$ ）
与えられるグラフは木

出力

答えとなる数値を1行で出力すること。

解法

単純に、すべての考えられうる $l, r$ に対して、頂点 $1$ から辿り着ける頂点を数え上げていく、という方法が考えられます。しかし、 $l, r$ をforループ等で決め、 $l, r$ を決めた後にも辺の本数が $n - 1$ 本あることより、この方法では $O({n}^{3})$ 程度の計算量がかかってしまうと考えられます。

この問題では最終的には「すべての $(l, r)$ の組に対して、頂点 $1$ から $l$ 以上 $r$ 以下の番号の辺を通って辿り着ける頂点の個数の総数」を求めます。これは、「ある頂点 $i$ に対して、頂点 $1$ から $l$ 以上 $r$ 以下の番号の辺を通って辿り着けるような、 $(l, r)$ の個数の総数」と等しいです。

すなわち、「辺の番号をもとにして、辿り着ける頂点の個数の総数」を求める問題を、「頂点をもとにして、辺の番号として満たされるものの個数の総数」を求める問題として言い換えることができます。

また、今回の問題で与えられるグラフは木構造であるため、頂点 $1$ から頂点 $i$ に辿り着く方法は一意に定まります。

ここで、頂点 $1$ でないある頂点 $i$ に対して、辿り着くために通る辺の番号の最小値を ${a}_{i}$ 、最大値を ${b}_{i}$ とします。（ $i$ は $1$ ではないので、 $2 \leq i \leq n$ です。）

頂点 $i$ に辿り着く方法が一意に定まることから、この ${a}_{i}, {b}_{i}$ は一意に定まります。

従って、問題の $(l, r)$ に頂点 $i$ が含まれるための条件は $1 \leq l \leq {a}_{i}$ かつ ${b}_{i} \leq r \leq (n - 1)$ となります。

以上から、頂点 $i$ が含まれるような $(l, r)$ の組み合わせの総数は ${a}_{i} (n - {b}_{i})$ であると言えます。

ここで、この ${a}_{i}, {b}_{i}$ については、頂点 $1$ から与えられたグラフ通りに辺を辿っていくことにより、求めることができます。

具体的には、既に頂点 $i$ に到達しているとして、辺 $k$ によって頂点 $j$ に辿り着けるとき、

${a}_{j} = \min ({a}_{i}, k)$
${b}_{j} = \max ({b}_{i}, k)$

と更新していくことができます。

また、頂点 $1$ については、必ず頂点 $1$ から出発することから、辺を辿らずとも必ず到達することができます。そのため、どのような $(l, r)$ に対しても必ず頂点 $1$ は含まれます。

$1 \leq l \leq r \leq n - 1$ であるため、 $r = i$ のときに $l$ の個数は $i$ 個となることから、 $(l, r)$ の総数は \begin{align} \sum^{n - 1}_{i = 1} i = \frac{1}{2} n(n - 1) \end{align} であると言えます。

以上から、求める個数については、頂点 $1$ の場合とそれ以外の場合を考えて、 \begin{align} \frac{1}{2} n (n - 1) + \sum^{n}_{i = 2} {a}_{i} (n - {b}_{i}) \end{align} となります。

計算量については、各頂点の ${a}_{i}, {b}_{i}$ の値を求める際に、辺を順番に辿っていく箇所に $O(n)$ が必要になります。従って、これは実行時間制限に十分に間に合うと言えます。

ソースコード

まず、頂点 $1$ 以外の頂点について、値 ${a}_{i}, {b}_{i}$ を求める部分を関数 dfs() として以下のように実装しました。

vector<Pil> graph[int(2e5 + 5)]; // graph[i] は {次の頂点, 辺の番号} が入っている pair<int, ll>
Pll range[int(2e5 + 5)]; // range[i] は {a[i], b[i]} の値が入っている pair<ll, ll>
bool visited[int(2e5 + 5)]; // すでにその頂点を辿ったかどうかを表す bool 値

void dfs(int start) { // start は始点となる頂点が入る（この問題では常に start = 1）
  visited[start] = true; // 始点は探索済とする
  queue<int> que; // 次に探索を行う頂点の入る queue
  que.push(start);

  while (!que.empty()) {
    int now = que.front();
    que.pop();

    for (Pil next : graph[now]) {
      int next_node = next.first; // 次に辿るであろう頂点の番号
      ll next_edge = next.second; // その頂点に辿り着くことのできる辺の番号

      if (visited[next_node]) { // もし次の頂点が探索済なら、このループは行わない
        continue;
      }

      visited[next_node] = true; // 次の頂点を探索済とする
      que.push(next_node); // 次の頂点を queue に入れる

      range[next_node].first = min(range[now].first, next_edge); // a[i] の値を next_edge を用いて更新
      range[next_node].second = max(range[now].second, next_edge); // b[i] の値を next_edge を用いて更新
    }
  }

  return;
}

以上の dfs() が実装されている状態で、main() 関数の中に答えを出力する部分を以下のように実装しました。

ll n;
vector<Pil> graph[int(2e5 + 5)];
Pll range[int(2e5 + 5)];

int main() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n - 1; i++) {
    int u, v;
    cin >> u >> v;

    // グラフに値を入れていく
    graph[u].push_back({v, i});
    graph[v].push_back({u, i});
  }

  range[1] = {n - 1, 1}; // dfs 用に、仮で a[1] = n - 1, b[1] = 1 としておく（後の計算には使用しない）
  dfs(1); // dfs を行う

  ll ans = (n - 1) * n / 2; // 頂点 1 の分を予め計算する
  for (int i = 2; i <= n; i++) {
    ans += range[i].first * (n - range[i].second); // 頂点 2 以降の部分について ans に加えていく
  }
  cout << ans << endl; // 答えの出力

  return 0;
}