Leaf Color

AtCoder

IDabc340_g

Time2000ms

Memory256MB

Difficulty—

There is a tree $T$ with $N$ vertices numbered from $1$ to $N$. The $i$\-th edge connects vertices $u_i$ and $v_i$. Additionally, vertex $i$ is painted with color $A_i$. Find the number, modulo $998244353$, of (non-empty) subsets $S$ of the vertex set of $T$ that satisfy the following condition: * The induced subgraph $G$ of $T$ by $S$ satisfies all of the following conditions: * $G$ is a tree. * All vertices with degree $1$ have the same color. What is an induced subgraph? Let $S$ be a subset of the vertices of a graph $G$. The induced subgraph of $G$ by $S$ is a graph whose vertex set is $S$ and whose edge set consists of all edges in $G$ that have both endpoints in $S$. ## Constraints * $1 \leq N \leq 2 \times 10^5$ * $1 \leq A_i \leq N$ * $1 \leq u_i \lt v_i \leq N$ * The graph given in the input is a tree. * All input values are integers. ## Input The input is given from Standard Input in the following format: $N$ $A_1$ $A_2$ $\dots$ $A_N$ $u_1$ $v_1$ $u_2$ $v_2$ $\vdots$ $u_{N-1}$ $v_{N-1}$ [samples]

Samples

Input #1

Output #1

4

The following four sets of vertices satisfy the condition.

*   $\lbrace 1 \rbrace$
*   $\lbrace 1, 2, 3 \rbrace$
*   $\lbrace 2 \rbrace$
*   $\lbrace 3 \rbrace$

Input #2

Output #2

Input #3

15
5 3 5 1 1 4 4 4 2 5 5 4 4 2 5
3 13
4 10
7 11
8 9
2 10
2 14
5 11
5 6
6 13
12 13
9 14
9 13
1 13
1 15

Output #3

API Response (JSON)

{
  "problem": {
    "name": "Leaf Color",
    "description": {
      "content": "There is a tree $T$ with $N$ vertices numbered from $1$ to $N$. The $i$\\-th edge connects vertices $u_i$ and $v_i$. Additionally, vertex $i$ is painted with color $A_i$.   Find the number, modulo $998",
      "description_type": "Markdown"
    },
    "platform": "AtCoder",
    "limit": {
      "time_limit": 2000,
      "memory_limit": 262144
    },
    "difficulty": "None",
    "is_remote": true,
    "is_sync": true,
    "sync_url": null,
    "sign": "abc340_g"
  },
  "statements": [
    {
      "statement_type": "Markdown",
      "content": "There is a tree $T$ with $N$ vertices numbered from $1$ to $N$. The $i$\\-th edge connects vertices $u_i$ and $v_i$. Additionally, vertex $i$ is painted with color $A_i$.  \nFind the number, modulo $998...",
      "is_translate": false,
      "language": "English"
    }
  ]
}

Full JSON Raw Segments