Anton は最近、EJOI の期間中に Kaunas を歩き回っていて、参加者がたむろする $N$ 個の地点を発見しました。地点には $0$ から $N-1$ までの番号が付いています。各地点はちょうど 1 つのチームのメンバーによって占められており、チームには $0$ から $K-1$ までの番号が付いています。同じチームのメンバーが任意の数の地点を占めていても構いません。また、どの地点も占めていないチームが存在しても構いません。
地点たちは $N-1$ 本の双方向道路で結ばれており、任意の 2 地点の間には単純パスがちょうど 1 つ存在します。すなわち、それらは木をなします。ここで、2 地点間の単純パスとは、相異なる地点の列であって、どの 2 つの連続する地点も道路で結ばれているものをいいます。パスの長さは、それが使う道路の本数、すなわち訪れる地点の数より 1 少ない値です。
Anton は、できるだけ多くの地点を訪れるように、単純パスに沿って散歩したいと思っています。単純パスは、その長さが木のすべての単純パスの中で最大であるとき、「興味深い」（interesting）と呼ばれます。パスの teamfulness とは、そのパスに沿って Anton が出会う異なるチームの数です。
あなたの仕事は、すべての異なる興味深いパスに対する teamfulness の合計を求めることです。2 つの興味深いパスは、それらがちょうど同じ地点集合を訪れる場合に限り同じものとみなされます（特に、逆向きにたどっても別のパスにはなりません）。
実装の詳細
あなたは以下の関数を実装する必要があります：
long long teamfulness(int N, int K, std::vector<int> a,
std::vector<int> u, std::vector<int> v)- $N$: 地点の数
- $K$: チームの数
- $a$: 長さ $N$ の整数列。$a_i$ は地点 $i$ を占めているチーム（各 $0 \le i < N$）
- $u, v$: 長さ $N-1$ の整数列。$u_i$ と $v_i$ は $i$ 番目の道路で結ばれている 2 地点（各 $0 \le i < N-1$）
この関数は各テストケースにつきちょうど 1 回呼び出され、すべての興味深いパスに対する teamfulness の合計を返さなければなりません。
制約
- $3 \le N \le 10^6$
- $1 \le K < N$
- $0 \le a_i < K$（各 $0 \le i < N$）
- $0 \le u_i, v_i < N$（各 $0 \le i < N-1$）
入出力例
入力 1
6 3 1 0 0 1 2 1 0 1 0 2 0 3 0 4 0 5
出力 1
21
注記 1
単純パスの最大長は 2（道路 2 本、地点 3 つ）なので、興味深いパスの長さは 2 です。teamfulness が 3 の興味深いパスが 2 本、teamfulness が 2 の興味深いパスが 7 本、teamfulness が 1 の興味深いパスが 1 本あり、合計は 21 になります。
入力 2
7 1 0 0 0 0 0 0 0 0 1 0 2 1 3 1 4 2 5 2 6
出力 2
4
注記 2
地点とそれらを結ぶ道路は以下の図のとおりです（チーム 0 は黄色で示されています）：
すべての地点にたむろしているチームはチーム 0 の 1 つだけなので、すべてのパスの teamfulness は 1 です。興味深いパス（長さ 4）は 4 本あり、したがってすべての興味深いパスに対する teamfulness の合計も 4 になります。
入力 3
6 3 0 1 2 0 1 2 0 1 1 2 2 3 1 4 2 5
出力 3
11
注記 3
地点とそれらを結ぶ道路は以下の図のとおりです（チーム 0 は黄色、チーム 1 は緑、チーム 2 は赤で示されています）：
興味深いパスの長さは 3 です。興味深いパスは全部で 4 本あり、そのうち 3 本の teamfulness は 3、1 本の teamfulness は 2 です。したがって、すべての興味深いパスに対する teamfulness の合計は 11 になります。
小課題
| 小課題 | 得点 | $N$ | $K$ | 追加の制約 |
|---|---|---|---|---|
| 0 | 0 | - | - | サンプル |
| 1 | 4 | $\le 10^6$ | $| すべての地点は他の高々 2 つの地点と直接結ばれている。 |
|
| 2 | 7 | $\le 10^6$ | $| ある 1 つの地点が他のすべての地点と直接結ばれている。 |
|
| 3 | 9 | $< 200$ | $|
| |
| 4 | 10 | $< 2 \cdot 10^3$ | ||
| 5 | 10 | $\le 10^6$ | $=1$ | |
| 6 | 9 | $\le 10^6$ | $<2$ | |
| 7 | 11 | $< 50$ | $|
| |
| 8 | 12 | $< 2 \cdot 10^5$ | $|
| |
| 9 | 13 | $\le 10^6$ | $| 興味深いパスの長さは奇数である。 |
|
| 10 | 15 | $< 10^6$ | $|
| |
入力
入力形式は以下の通りです：
- 1 行目: 2 つの整数 – $N$ と $K$ の値
- 2 行目: $N$ 個の整数 $a_0, a_1, \dots, a_{N-1}$。ここで $a_i$ は地点 $i$ を占めているチーム
- $3+i$ 行目: 2 つの整数 $u$ と $v$ – $i$ 番目の道路で結ばれている 2 地点
出力
出力形式は以下の通りです：
- 1 行目: 1 つの整数 – 関数呼び出しの戻り値