Problem p02749 (python)

Fixed submission:
from collections import defaultdict
import sys
input = sys.stdin.readline


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

colors = [-1] * n

colors[0] = 0

todo = []
for to in edges[0]:
    todo.append([to, 1])

while todo:
    node, color = todo.pop()

    if colors[node] != -1:
        continue

    colors[node] = color
    new_color = 1 - color
    for to in edges[node]:
        if colors[to] != -1:
            continue
        todo.append([to, new_color])

red_count = sum(colors)  # 1
blue_count = n - red_count  # 0

ans = [0] * n

num3s = list(range(3, n+1, 3))
nums = [num for num in range(1, n+1) if num % 3 != 0]

if red_count <= n // 3:
    for node in range(n):
        if colors[node] == 1:
            ans[node] = num3s.pop()

        else:
            if nums:
                num = nums.pop()
            else:
                num = num3s.pop()
            ans[node] = num


elif blue_count <= n // 3:
    for node in range(n):
        if colors[node] == 0:
            ans[node] = num3s.pop()

        else:
            if nums:
                num = nums.pop()
            else:
                num = num3s.pop()
            ans[node] = num

else:
    num1s = list(range(1, n+1, 3))
    num2s = list(range(2, n+1, 3))

    for node in range(n):
        if colors[node] == 1:
            if num1s:
                num = num1s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

        else:
            if num2s:
                num = num2s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

print(*ans)

Buggy submission (verdict: Runtime Error):
from collections import defaultdict


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

colors = [-1] * n

colors[0] = 0

todo = []
for to in edges[0]:
    todo.append([to, 1])

while todo:
    node, color = todo.pop()

    if colors[node] != -1:
        continue

    colors[node] = color
    new_color = 1 - color
    for to in edges[node]:
        if colors[to] != -1:
            continue
        todo.append([to, new_color])

red_count = sum(colors)  # 1
blue_count = n - red_count  # 0

ans = [0] * n

num3s = list(range(3, n+1, 3))
nums = [num for num in range(1, n+1) if num % 3 != 0]

if red_count <= n // 3:
    for node in range(n):
        if colors[node] == 1:
            ans[node] = num3s.pop()

        else:
            if num3s:
                num = num3s.pop()
            else:
                num = nums.pop()
            ans[node] = num


elif blue_count <= n // 3:
    for node in range(n):
        if colors[node] == 0:
            ans[node] = num3s.pop()

        else:
            if num3s:
                num = num3s.pop()
            else:
                num = nums.pop()
            ans[node] = num

else:
    num1s = list(range(1, n+1, 3))
    num2s = list(range(2, n+1, 3))
    print(1)
    exit()

    for node in range(n):
        if colors[node] == 1:
            if num1s:
                num = num1s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

        else:
            if num2s:
                num = num2s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

print(*ans)

Buggy submission (verdict: Runtime Error):
from collections import defaultdict


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

colors = [-1] * n

colors[0] = 0

todo = []
for to in edges[0]:
    todo.append([to, 1])

while todo:
    node, color = todo.pop()

    if colors[node] != -1:
        continue

    colors[node] = color
    new_color = 1 - color
    for to in edges[node]:
        if colors[to] != -1:
            continue
        todo.append([to, new_color])

red_count = sum(colors)  # 1
blue_count = n - red_count  # 0

ans = [0] * n

num3s = list(range(3 ,n+1, 3))
nums = [num for num in range(1, n+1) if num % 3 != 0]

if red_count <= n // 3:
    for node in range(n):
        if colors[node] == 1:
            ans[node] = num3s.pop()

        else:
            if num3s:
                num = num3s.pop()
            else:
                num  = nums.pop()
            ans[node] = num


elif blue_count <= n // 3:
    for node in range(n):
        if colors[node] == 0:
            ans[node] = num3s.pop()

        else:
            if num3s:
                num = num3s.pop()
            else:
                num = nums.pop()
            ans[node] = num

else:
    num1s = list(range(1, n+1, 3))
    num2s = list(range(2, n+1, 3))

    for node in range(n):
        if colors[node] == 1:
            if num1s:
                num = num1s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

        else:
            if num2s:
                num = num2s.pop()
            else:
                num = num3s.pop()
            ans[node] = num

print(*ans)