Submission #1239762

#	Time	Username	Problem	Language	Result	Execution time	Memory
1239762		woe	Team Coding (EGOI24_teamcoding)	Pypy 3	0 / 100	139 ms	48808 KiB

Main

from collections import defaultdict, deque

def solve():
    import sys
    input = sys.stdin.read
    data = input().split()

    idx = 0
    N = int(data[idx])
    idx += 1
    K = int(data[idx])
    idx += 1

    lang = list(map(int, data[idx:idx+N]))
    idx += N

    tree = [[] for _ in range(N)]
    parent = [-1]*N
    for i in range(1, N):
        b = int(data[idx])
        idx += 1
        tree[b].append(i)
        parent[i] = b

    depth = [0]*N
    def dfs_depth(node, d):
        depth[node] = d
        for child in tree[node]:
            dfs_depth(child, d+1)
    dfs_depth(0, 0)

    # Precompute depth-wise counts of languages
    depth_lang_count = defaultdict(lambda: defaultdict(int))
    for i in range(N):
        depth_lang_count[depth[i]][lang[i]] += 1

    def dfs_subtree(node):
        count = defaultdict(lambda: defaultdict(int))  # depth -> lang -> count
        q = deque([(node, depth[node])])
        while q:
            u, d = q.popleft()
            count[d][lang[u]] += 1
            for v in tree[u]:
                q.append((v, depth[v]))
        return count

    max_team = 0
    min_switches = 0

    for lead in range(N):
        lead_lang = lang[lead]
        subtree_counts = dfs_subtree(lead)

        team_size = 0
        switch_needed = 0

        for d in subtree_counts:
            same_in = subtree_counts[d][lead_lang]
            same_out = depth_lang_count[d][lead_lang] - same_in
            diff_in = sum(subtree_counts[d].values()) - same_in
            swappable = min(diff_in, same_out)

            team_size += same_in + swappable
            switch_needed += swappable

        if team_size > max_team or (team_size == max_team and switch_needed < min_switches):
            max_team = team_size
            min_switches = switch_needed

    print(max_team, min_switches)

Compilation message (stdout)

Compiling 'Main.py'...

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

=======

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