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



Sample input:
4
1 2
1 3
1 4
1100


Sample output:
4

Submission status: Wrong Answer

Solution code:
def solve():
    n = int(input())
    g = [[] for _ in range(n)]

    for _ in range(n - 1):
        a, b = [int(x) for x in input().split()]
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    ds = dict()

    parent = [-1] * n
    dfs_order = [0]
    for i in range(n):
        a = dfs_order[i]
        for b in g[a]:
            if b == parent[a]:
                continue
            dfs_order.append(b)
            parent[b] = a

    vd = [-1] * n

    for a in reversed(dfs_order):
        vd[a] = 0
        for b in g[a]:
            if b == parent[a]:
                continue
            ds[(a, b)] = vd[b]
            vd[a] = max(vd[a], vd[b])

        vd[a] = vd[a] + 1

    ans = 0
    d_parent = [0] * n

    for a in dfs_order:
        cd = [(d_parent[a], parent[a])]
        for b in g[a]:
            if b == parent[a]:
                continue
            cd.append((ds[(a, b)], b))

        cd.sort()
        cd.reverse()

        cur = (cd[1][0] + 1) if (len(cd) > 1) else 1
        cur = min(cur, cd[0][0] - 1)
        ans += cur + 1

        for b in g[a]:
            if b == parent[a]:
                continue
            x = cd[0][0]
            if cd[0][1] == b:
                x = cd[1][0]
            d_parent[b] = x + 1

    print(ans + 1)


if __name__ == '__main__':
    # t = int(input())
    t=1
    for _ in range(t):
        solve()

Submission status: Accepted

Solution code:
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()

Submission status: Runtime Error

Solution code:


def solve():
    n = int(input())
    g = [[] for _ in range(n)]

    for _ in range(n - 1):
        a, b = [int(x) for x in input().split()]
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    ds = dict()

    def dfs1(g, a, fr):
        res = 0
        for b in g[a]:
            if b == fr:
                continue
            x = dfs1(g, b, a)
            ds[(a, b)] = x
            res = max(res, x)

        return res + 1

    dfs1(g, 0, -1)
    # print(ds)
    ans = 0

    def dfs2(g, a, fr, d_fr):
        nonlocal ans
        # print(a, fr, d_fr)
        cd = [(d_fr, fr)]
        for b in g[a]:
            if b == fr:
                continue
            cd.append((ds[(a, b)], b))

        cd.sort()
        cd.reverse()
        # print(cd)
        cur = (cd[1][0] + 1) if (len(cd) > 1) else 1
        cur = min(cur, cd[0][0] - 1)
        # print(cur)
        ans += cur + 1

        for b in g[a]:
            if b == fr:
                continue
            x = cd[0][0]
            if cd[0][1] == b:
                x = cd[1][0]
            dfs2(g, b, a, x + 1)

    dfs2(g, 0, -1, 0)
    print(ans + 1)


if __name__ == '__main__':
    # t = int(input())
    t = 1
    for _ in range(t):
        solve()

Submission status: Runtime Error

Solution code:


def solve():
    n = int(input())
    g = [[] for _ in range(n)]

    for _ in range(n - 1):
        a, b = [int(x) for x in input().split()]
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    ds = dict()

    def dfs1(g, a, fr):
        res = 0
        for b in g[a]:
            if b == fr:
                continue
            x = dfs1(g, b, a)
            ds[(a, b)] = x
            res = max(res, x)

        return res + 1

    dfs1(g, 0, -1)
    # print(ds)
    ans = 0

    def dfs2(g, a, fr, d_fr):
        nonlocal ans
        # print(a, fr, d_fr)
        cd = [(d_fr, fr)]
        for b in g[a]:
            if b == fr:
                continue
            cd.append((ds[(a, b)], b))

        cd.sort()
        cd.reverse()
        # print(cd)
        cur = (cd[1][0] + 1) if (len(cd) > 1) else 1
        cur = min(cur, cd[0][0] - 1)
        # print(cur)
        ans += cur + 1

        for b in g[a]:
            if b == fr:
                continue
            x = cd[0][0]
            if cd[0][1] == b:
                x = cd[1][0]
            dfs2(g, b, a, x + 1)

    dfs2(g, 0, -1, 0)
    print(ans + 1)


if __name__ == '__main__':
    t = int(input())
    for _ in range(t):
        solve()