Problem:
 Score : 1900 points    Problem Statement There is a tree with N vertices.
The vertices are numbered 1 through N .
For each 1 ≤ i ≤ N - 1 , the i -th edge connects vertices a_i and b_i .
The lengths of all the edges are 1 .  Snuke likes some of the vertices.
The information on his favorite vertices are given to you as a string s of length N .
For each 1 ≤ i ≤ N , s_i is 1 if Snuke likes vertex i , and 0 if he does not like vertex i .  Initially, all the vertices are white.
Snuke will perform the following operation exactly once:   Select a vertex v that he likes, and a non-negative integer d . Then, paint all the vertices black whose distances from v are at most d .   Find the number of the possible combinations of colors of the vertices after the operation.      Constraints  2 ≤ N ≤ 2×10^5  1 ≤ a_i, b_i ≤ N  The given graph is a tree.  |s| = N  s consists of 0 and 1 .  s contains at least one occurrence of 1 .       Partial Score  In the test set worth 1300 points, s consists only of 1 .         Input The input is given from Standard Input in the following format:  N  a_1  b_1  a_2  b_2  :  a_{N - 1}  b_{N - 1}  s       Output Print the number of the possible combinations of colors of the vertices after the operation.        Sample Input 1 4
1 2
1 3
1 4
1100      Sample Output 1 4  The following four combinations of colors of the vertices are possible:          Sample Input 2 5
1 2
1 3
1 4
4 5
11111      Sample Output 2 11  This case satisfies the additional constraint for the partial score.       Sample Input 3 6
1 2
1 3
1 4
2 5
2 6
100011      Sample Output 3 8

Solution (Python):
import sys
import json
import itertools

INFINITY = 100000000

def parse_data(input_string):
    INF = INFINITY

    first_nl = input_string.index('\n')
    n = int(input_string[:first_nl])
    fav_pattern = [int(s) * INF for s in input_string[-n - 1:-1]]

    edges = json.loads("[" + input_string[first_nl + 1:-n - 2].replace(" ", ",").replace("\n", ",") + "]")
    graph = [[] for _ in range(n)]
    graph[0].append(None)
    for i in range(0, (n - 1) * 2, 2):
        a = edges[i] - 1
        b = edges[i + 1] - 1
        graph[a].append(b)
        graph[b].append(a)
    return (n, fav_pattern, graph)

def solve():
    INF = INFINITY

    n, fav_count, graph = parse_data(sys.stdin.read())
    height1 = [0] * n
    height2 = [0] * n
    dfs_order = []

    stack = [None, 0, 0]
    while True:
        node = stack[-1]
        index = stack[-2]
        if index < len(graph[node]):
            dtr = graph[node][index]
            stack[-2] += 1
            if dtr != stack[-3]:
                stack.extend((0, dtr))
                dfs_order.append(dtr)
            else:
                graph[node][0], graph[node][index] = dtr, graph[node][0]
        else:
            del stack[-2:]
            if len(stack) <= 1: break
            parent = stack[-1]
            dtr_height = height1[node] + 1
            if dtr_height >= height1[parent]:
                height1[parent], height2[parent] = dtr_height, height1[parent]
            elif dtr_height > height2[parent]:
                height2[parent] = dtr_height
            if fav_count[node] != 0:
                fav_count[parent] += 1

    answer = 0
    if fav_count[0] >= INF:
        lb = -1
    else:
        lb = min((height1[dtr] for dtr in graph[0][1:] if fav_count[dtr] != 0), default=INF)
    ub = height2[0]
    if ub > lb:
        answer += ub - lb
    for node in dfs_order:
        parent = graph[node][0]
        has_fav_below = (fav_count[node] != 0)
        has_fav_above = fav_count[parent] > has_fav_below
        fav_count[node] += has_fav_above
        node_height = height1[node]
        if node_height + 1 == height1[parent]:
            lb = height2[parent]
            if lb <= node_height and has_fav_above:
                answer += 1
            elif lb >= node_height and has_fav_below:
                answer += 1

            parent_height = lb + 1
            if parent_height >= node_height:
                height1[node], height2[node] = parent_height, node_height
            elif parent_height > height2[node]:
                height2[node] = parent_height
        else:
            answer += has_fav_below
            lb = height1[parent]
            height1[node], height2[node] = lb + 1, node_height
        if fav_count[node] >= INF:
            lb = -1
        else:
            if not has_fav_above: lb = INF
            for dtr in itertools.islice(graph[node], 1, None):
                if fav_count[dtr] != 0:
                    lb = min(height1[dtr], lb)
        ub = height2[node]
        if ub > lb:
            answer += ub - lb

    print(answer)

solve()