Submission #1190666

#TimeUsernameProblemLanguageResultExecution timeMemory
1190666ofozTeam Coding (EGOI24_teamcoding)Pypy 3
0 / 100
4006 ms1114112 KiB
from math import floor, ceil
from sys import stdin, setrecursionlimit

setrecursionlimit(int(2e5))
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()

Compilation message (stdout)

Compiling 'Main.py'...

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

=======
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...