Werewolves

Time limit2sMemory limit1024 MB

Summary

Count connected subgraphs of a colored tree in which one color makes up strictly more than half of the nodes, modulo 998244353.

Level: Hard8 of 10
Topics: Tree, Divide and conquer, Dynamic programming
Solved: No attempts yet

Problem

You are given a colored tree on $n$ nodes, and node $i$ has color $c_i$ . A tree on $n$ nodes is a connected graph with $n-1$ edges.

Count the connected subgraphs of this tree for which some color (the majority color) appears on strictly more than half of the subgraph's nodes. Since the answer can be large, print it modulo $998\,244\,353$ .

Input

The first line contains one integer $n$ ( $1 \le n \le 3000$ ), the number of nodes. The second line contains $n$ integers $c_1, c_2, \ldots, c_n$ ( $1 \le c_i \le n$ ), the colors of the nodes. Each of the next $n-1$ lines contains two integers $u_i$ and $v_i$ ( $1 \le u_i, v_i \le n$ , $u_i \neq v_i$ ), an edge of the tree. The given graph is a tree.

Output

Print one integer, the number of connected subgraphs with a majority color, modulo $998\,244\,353$ .

Hint

In the figures, color 1 is blue, color 2 is red, and color 3 is yellow. In the first figure, the subgraph on nodes 1 and 2 has no majority color because red and yellow each appear once.

Examples2

Example 1
Input
```
3
2 3 3
1 2
2 3
```
Expected output
```
5
```
Example 2
Input
```
4
1 1 3 3
1 2
1 3
1 4
```
Expected output
```
8
```