Có \(N\) địa điểm tạo thành một cây. Địa điểm \(i\) thuộc đội \(a_i\in[0,K-1]\); một đội có thể chiếm nhiều địa điểm hoặc không có địa điểm nào.

Một đường đi đơn được gọi là thú vị nếu có độ dài lớn nhất trong mọi đường đi đơn của cây, tức là một đường kính. Độ đa dạng đội của đường đi là số đội phân biệt xuất hiện trên đó.

Hãy tính tổng độ đa dạng đội trên mọi đường đi thú vị khác nhau. Hai đường đi được xem là giống nhau khi chúng đi qua cùng tập đỉnh, nên hai hướng đi không tạo hai đường khác nhau.

Giao diện thư viện

Submission C++ phải include teamfulness.h và cài đặt:

long long teamfulness(
    int N, int K,
    std::vector<int> a,
    std::vector<int> u,
    std::vector<int> v);

Hai vector u,v có \(N-1\) phần tử và cạnh thứ \(i\) nối u[i] với v[i]. Hàm được gọi đúng một lần.

Dữ liệu vào

Submission không đọc standard input. Sample grader đọc \(N,K\), dãy đội và \(N-1\) cạnh.

Dữ liệu ra

Submission không ghi standard output. Kết quả được trả từ teamfulness.

Ràng buộc

• \(3\le N\le10^6\).
• \(1\le K\le N\).
• \(0\le a_i<K\).
• \(0\le u_i,v_i<N\).
• Các cạnh tạo thành một cây.

Phân nhóm

1. \(4\) điểm: bậc mỗi đỉnh không quá \(2\).
2. \(7\) điểm: có một đỉnh kề trực tiếp với mọi đỉnh còn lại.
3. \(9\) điểm: \(N\le200\).
4. \(10\) điểm: \(N\le2000\).
5. \(10\) điểm: \(K=1\).
6. \(9\) điểm: \(K\le2\).
7. \(11\) điểm: \(N\le2\cdot10^5\), \(K\le50\).
8. \(12\) điểm: \(N\le2\cdot10^5\).
9. \(13\) điểm: độ dài đường đi thú vị là số lẻ.
10. \(15\) điểm: không có ràng buộc thêm.

Ví dụ

6 3
1 0 0 1 2 1
0 1
0 2
0 3
0 4
0 5

21

7 1
0 0 0 0 0 0 0
0 1
0 2
1 3
1 4
2 5
2 6

4

6 3
0 1 2 0 1 2
0 1
1 2
2 3
1 4
2 5

11

Nguồn

EJOI 2026 - Ngày 2, Teamfulness.

Đề bài EJOI 2026 được phát hành theo giấy phép Creative Commons Attribution (CC BY).