Олимпиадный тренинг

Задача . F1. Разрезание дерева (легкая версия)


Задано неориентированное дерево из \(n\) вершин.

Некоторые вершины покрашены в красный цвет, некоторые — в синий, а некоторые не покрашены совсем. Гарантируется, что дерево содержит хотя бы одну красную вершину и хотя бы одну синюю вершину.

Вы выбираете ребро и удаляете его из дерева. Дерево распадается на две связных компоненты. Назовем ребро хорошим, если в каждой из двух компонент не будет одновременно синей и красной вершин.

Сколько хороших ребер в данном дереве?

Входные данные

В первой строке записано одно целое число \(n\) (\(2 \le n \le 3 \cdot 10^5\)) — количество вершин в дереве.

Во второй строке записаны \(n\) целых чисел \(a_1, a_2, \dots, a_n\) (\(0 \le a_i \le 2\)) — цвета вершин. \(a_i = 1\) значит, что вершина \(i\) покрашена в красный, \(a_i = 2\) значит, что вершина \(i\) покрашена в синий, и \(a_i = 0\) значит, что вершина \(i\) не покрашена.

В \(i\)-й из следующих \(n - 1\) строк содержится два целых числа \(v_i\) и \(u_i\) (\(1 \le v_i, u_i \le n\), \(v_i \ne u_i\)) — ребра дерева. Гарантируется, что данные ребра образуют дерево. Гарантируется, что дерево содержит хотя бы одну красную вершину и хотя бы одну синюю вершину.

Выходные данные

Выведите одно целое число — количество хороших ребер в данном дереве.

Примечание

Дерево из первого примера:

Единственное хорошее ребро — это ребро \((2, 4)\). После его удаления дерево распадается на компоненты \(\{4\}\) и \(\{1, 2, 3, 5\}\). В первой компоненте есть только красная вершина, а во второй — только синие и непокрашенные вершины.

Дерево из второго примера:

Каждое ребро в нем хорошее.

Дерево из третьего примера:

Ребро \((1, 3)\) разделяет дерево на компоненты \(\{1\}\) и \(\{3, 2\}\), во второй есть одновременно красная и синяя вершины, поэтому это ребро не хорошее. Ребро \((2, 3)\) разделяет дерево на компоненты \(\{1, 3\}\) и \(\{2\}\), в первой есть одновременно красная и синяя вершины, поэтому это ребро также не хорошее. Потому ответ равен 0.


Примеры
Входные данныеВыходные данные
1 5
2 0 0 1 2
1 2
2 3
2 4
2 5
1
2 5
1 0 0 0 2
1 2
2 3
3 4
4 5
4
3 3
1 1 2
2 3
1 3
0

time 2000 ms
memory 256 Mb
Правила оформления программ и список ошибок при автоматической проверке задач

Статистика успешных решений по компиляторам
 Кол-во
С++ Mingw-w645
Комментарий учителя