東京農工大学MCCプログラミングコンテスト2023Winter
コンテスト日時
2024/03/25 (Mo) 20:00 - 22:00

E - Destroy Tree

Benihuki
2
s
1024
MB
100

問題文

頂点 $1, 2, \dots, N$ からなる $N$ 頂点の根付き木 $T$ が与えられます。

頂点 $1$ はこの木の根です。 $i = 1, 2, \dots, N − 1$ について、 $i$ 番目の辺は頂点 $A_i$ と頂点 $B_i$ を結んでいます。

また、 $(2, 3, \dots, N)$ を並び替えて得られる数列 $x$ が与えられます。ここで、 $x$ の $i$ 番目の項を $x_i$ とします。

あなたは木 $T$ と数列 $x$ を用いて次の操作を $N - 1$ 回行います。はじめ、スコアは $0$ です。

$i$ 回目の操作では次の操作を行います。

  1. 頂点 $x_i$ がその時点における木 $T$ に存在しない場合、何もしない。
  2. 頂点 $x_i$ がその時点における木 $T$ に存在する場合、次の操作を行う。
    • 木 $T$ から頂点 $x_i$ を削除する。このとき、頂点 $x_i$ を根とする部分木に含まれる頂点であって、木 $T$ に残っているすべての頂点も同時に削除される。
    • この操作によって削除された頂点数を $v$ としたとき、 $v^2$ のスコアを得る。

すべての操作を終えた後の合計のスコアを求めてください。

制約

Hard (50点)

  • $2 \le N \le 2 \times 10^5$
  • $1 \le A_i < B_i \le N$
  • $T$ は木である
  • $2 \le x_i \le N$
  • $i \ne j$ ならば $x_i \ne x_j$
  • 入力はすべて整数

Normal (25点)

  • Hard の制約に以下の制約を追加
  • $2 \le N \le 3000$

Easy (25点)

  • Hard の制約に以下の制約を追加
  • $2 \le N \le 3000$
  • $i$ 本目の辺は頂点 $i$ と頂点 $i + 1$ を結ぶ

入力

入力は以下の形式で標準入力から与えられる。

$N$
$A_1 \ \ B_1$
$A_2 \ \ B_2$
$\ \ \vdots \quad \ \vdots$
$A_{N-1} \ \ B_{N-1}$
$x_1 \ \ x_2 \ \ \cdots \ \ x_{N-1}$

出力

答えを出力せよ。

入力例 1
5 1 2 1 3 2 4 2 5 4 2 5 3
出力例 1
6

操作は次のように行われます。

  • 頂点 $x_1 = 4$ は木 $T$ に残っているため、頂点 $4$ を削除する。頂点が $1$ つ削除されたので、スコア $1$ を得る。
  • 頂点 $x_2 = 2$ は木 $T$ に残っているため、頂点 $2$ を削除する。このとき、頂点 $5$ は頂点 $2$ を根とする部分木に含まれるため、頂点 $5$ も削除される。頂点が $2$ つ削除されたので、スコア $4$ を得る。
  • 頂点 $x_3 = 5$ は木 $T$ に残っていないため、何もしない。このときスコアは得られない。
  • 頂点 $x_4 = 3$ は木 $T$ に残っているため、頂点 $3$ を削除する。頂点が $1$ つ削除されたので、スコア $1$ を得る。

操作全体で $1 + 4 + 0 + 1 = 6$ のスコアを得たので、 $6$ を出力します。

tdg_1

この入力例は Normal・Hard の制約を満たします。

入力例 2
5 1 2 2 3 3 4 4 5 4 2 5 3
出力例 2
8

操作は次のように行われます。

  • 頂点 $x_1 = 4$ は木 $T$ に残っているため、頂点 $4$ を削除する。このとき、頂点 $5$ は頂点 $4$ を根とする部分木に含まれるため、頂点 $5$ も削除される。頂点が $2$ つ削除されたので、スコア $4$ を得る。
  • 頂点 $x_2 = 2$ は木 $T$ に残っているため、頂点 $2$ を削除する。このとき、頂点 $3$ は頂点 $2$ を根とする部分木に含まれるため、頂点 $3$ も削除される。頂点が $2$ つ削除されたので、スコア $4$ を得る。
  • 頂点 $x_3 = 5$ は木 $T$ に残っていないため、何もしない。このときスコアは得られない。
  • 頂点 $x_4 = 3$ は木 $T$ に残っていないため、何もしない。このときスコアは得られない。

操作全体で $4 + 4 + 0 + 0 = 8$ のスコアを得たので、 $8$ を出力します。

dt2

この入力例は Easy・Normal・Hard の制約を満たします。

入力例 3
8 1 2 2 3 2 4 3 5 3 6 4 7 4 8 2 8 7 6 5 4 3
出力例 3
49

この入力例は Normal・Hard の制約を満たします。

提出
C++23 (g++ 12.2.0)