제출 #1190663

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1190663		ofoz	Team Coding (EGOI24_teamcoding)	Pypy 3	0 / 100	654 ms	1114112 KiB

Main

from math import floor, ceil
from sys import stdin
def solve():
    n, k = map(int, input().split(" "))
    col = list(map(int, input().split(" ")))
    adj = [[] for _ in range(n)]
    for i in range(1, n):
        x = int(input())
        adj[x].append(i)

    def dfs1(v: int, level: int, vertices: list[int]): # Calculate the number of vertices in a given level of a subtree
        vertices[level] += 1
        for to in adj[v]: dfs1(to, level+1, vertices)

    def dfs2(v: int, level: int, colorCnt: list[list[int]]): # Calculate the number of vertices of color c at level i
        colorCnt[level][col[v]] += 1
        for to in adj[v]: dfs2(to, level+1, colorCnt)
    

    res = [0, 0]
    def dfs(v: int, level: int, colorCnt: list[list[int]]):
        nonlocal res
        c = col[v]
        vertices = [0] * n
        dfs1(v, level, vertices)
        colorCntSubtree = [[0] * k for _ in range(n)]
        dfs2(v, level, colorCntSubtree)

        coders = swaps = 0
        for l in range(level+1, n):
            swaps += min(vertices[l] - colorCntSubtree[l][c], colorCnt[l][c] - colorCntSubtree[l][c])
            # if v == 0 and swaps:
                # print(l, vertices[l] - colorCntSubtree[l][c])
                # return
            coders += min(vertices[l], colorCnt[l][c])
        
        if coders > res[0] or (coders == res[0] and swaps < res[1]):
            res = [coders, swaps]

        for to in adj[v]:
            dfs(to, level+1, colorCnt)
        


    colorCnt = [[0] * k for _ in range(n)] # colorCnt[l][k] denotes number of vertices of color k at level l
    dfs2(0, 0, colorCnt)

    dfs(0, 0, colorCnt)
    print(res[0]+1, res[1])



solve()

컴파일 시 표준 출력 (stdout) 메시지

Compiling 'Main.py'...

=======
  adding: __main__.pyc (deflated 52%)

=======

• Subtask 1: 0 / 12
• Subtask 2: 0 / 19
• Subtask 3: 0 / 27
• Subtask 4: 0 / 23
• Subtask 5: 0 / 19