Submission #1045438

#	Time	Username	Problem	Language	Result	Execution time	Memory
1045438		RauPro	Mergers (JOI19_mergers)	Pypy 3	0 / 100	285 ms	47652 KiB

mergers

import sys
from collections import defaultdict
from sys import setrecursionlimit, stdin, stdout

setrecursionlimit(200000)
class DSU:
    def __init__(self, n):
        self.par = list(range(n))

    def find(self, x):
        if self.par[x] == x:
            return x
        self.par[x] = self.find(self.par[x])
        return self.par[x]

    def merge(self, a, b):
        a = self.find(a)
        b = self.find(b)
        if a != b:
            self.par[b] = a


def setup(v, parent):
    for p in graf[v]:
        if p == parent:
            continue
        depth[p] = depth[v] + 1
        up[p] = v
        setup(p, v)


def merge(a, b):
    a = dsu.find(a)
    b = dsu.find(b)
    while a != b:
        if depth[a] < depth[b]:
            a, b = b, a
        dsu.merge(up[a], a)
        a = dsu.find(up[a])



n, k = map(int, input().split())

graf = [[] for i in range(n+1)]
dsu = DSU(n)
up = [0] * n
depth = [0] * n

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

setup(0, -1)

groups = [[] for i in range(n+1)]
for i in range(n):
    a = int(input()) - 1
    groups[a].append(i)

for group in groups.values():
    for j in range(1, len(group)):
        merge(group[0], group[j])

sol = 0
deg = [0] * n
for i in range(n):
    for p in graf[i]:
        if dsu.find(i) != dsu.find(p):
            if deg[dsu.find(i)] == 0:
                sol += 1
            if deg[dsu.find(i)] == 1:
                sol -= 1
            deg[dsu.find(i)] += 1

print((sol + 1) // 2)

• Subtask 1: 0 / 10
• Subtask 2: 0 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30