Submission #1048081

#	Time	Username	Problem	Language	Result	Execution time	Memory
1048081		RauPro	Mergers (JOI19_mergers)	Pypy 3	0 / 100	204 ms	53520 KiB

mergers

import os
import sys
from collections import *
from heapq import *
from math import gcd, floor, ceil, sqrt
from copy import deepcopy
from itertools import permutations, combinations, product
from bisect import bisect_left, bisect_right
from functools import lru_cache, reduce
import operator
from random import getrandbits

# Para mejorar el rendimiento de la entrada/salida
input = lambda: sys.stdin.readline().strip()
flush = lambda: sys.stdout.flush()
print = lambda *args, **kwargs: sys.stdout.write(' '.join(map(str, args)) + kwargs.get("end", "\n")) and flush()

# Optimización de la recursión para Python
sys.setrecursionlimit(100000)


# Funciones para lectura de múltiples tipos de datos
def ints(): return map(int, input().split())
def strs(): return input().split()
def chars(): return list(input().strip())
def mat(n): return [list(ints()) for _ in range(n)]  # Matriz de n x m donde m es el número de enteros en una línea


# Constantes útiles
INF = float('inf')
MOD = 1000000007  # Modulo por defecto, cambiar si se necesita otro
abcd = "abcdefghijklmnopqrstuvwxyz"

# Algunas funciones útiles
def add(x, y, mod=MOD): return (x + y) % mod
def sub(x, y, mod=MOD): return (x - y) % mod
def mul(x, y, mod=MOD): return (x * y) % mod

# Inverso multiplicativo de a modulo m (cuando m es primo)
def invmod(a, mod=MOD): return powmod(a, mod - 2, mod)

def lcm(a, b): return a * b // gcd(a, b)

RANDOM = getrandbits(32)

class Wrapper(int):
    def __init__(self, x):
        int.__init__(x)
    def __hash__(self):
        return super(Wrapper, self).__hash__() ^ RANDOM


depth = []
parent = []
visited = []
AL = []
uf = None


class UF:
    def __init__(self, n):
        self.parents = list(range(n))
        self.ranks = [0] * n
        self.sizes = [1] * n
        self.numdisjoint = n

    def find(self, x):
        xp = x
        children = []
        while xp != self.parents[xp]:
            children.append(xp)
            xp = self.parents[xp]
        for c in children:
            self.parents[c] = xp
        return xp

    def union(self, a, b):
        ap = self.find(a)
        bp = self.find(b)
        if ap == bp:
            return

        if self.ranks[ap] < self.ranks[bp]:
            self.parents[ap] = bp
            self.sizes[bp] += self.sizes[ap]
        elif self.ranks[bp] < self.ranks[ap]:
            self.parents[bp] = ap
            self.sizes[ap] += self.sizes[bp]
        else:
            self.parents[bp] = ap
            self.ranks[ap] += 1
            self.sizes[ap] += self.sizes[bp]

        self.numdisjoint -= 1

    def size(self, x):
        return self.sizes[self.find(x)]


# Tree compression from vertex a to b
def path_compression(a, b):
    a = uf.find(a)
    b = uf.find(b)
    while depth[a] < depth[b]:
        uf.union(a, b)
        b = parent[b]
        path_compression(a, b)


def dfs(u, lvl, p):
    global AL, visited, parent, depth, uf
    visited[u] = True
    depth[u] = lvl
    parent[u] = p
    for v in AL[u]:
        if not visited[v]:
            dfs(v, lvl+1, u)
def main():
    global AL, visited, parent, depth, uf
    n, m = ints()
    AL = [[] for i in range(n+1)]
    for i in range(n-1):
        u,v = ints()
        AL[u].append(v)
        AL[v].append(u)
    tree_groups = [[] for i in range(n+1)]
    for i in range(n):
        g = int(input())
        tree_groups[g].append(i+1)

    #print(tree_groups)
    visited = [False] * (n+1)
    parent = [0] * (n+1)
    depth = [0] * (n + 1)
    dfs(1, 0, 1)

    #print(depth, parent)
    uf = UF(n+1)
    for group in tree_groups:
        for i in range(1, len(group)):
            path_compression(group[0], group[i])

    ans = 0
    for u in range(1, n+1):
        if uf.size(u) == 1:
            ans +=1
    print((ans + 1) // 2)




if __name__ == "__main__":
    main()

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