Problem p02749 (python)

Fixed submission:
from sys import setrecursionlimit

setrecursionlimit(10**6)

def dfs(i, depth):
    visited[i] = True
    oe = depth%2
    ans[i-1] = oe
    count[oe] += 1
    nodes = tree[i]
    for node in nodes:
        if visited[node]:
            continue
        dfs(node, depth+1)
    


N = int(input())
ans = [-1]*N
tree = [[] for i in range(N+1)]
visited = [False]*(N+1)
count = [0, 0]
base = [1, 2, 3]
for i in range(N-1):
    a, b = map(int, input().split())
    tree[a] += [b]
    tree[b] += [a]
dfs(1, 0)
even_len = count[0]
odd_len = count[1]
if even_len > N//3 and odd_len > N//3:
    for i, oe in enumerate(ans):
        if base[oe] > N:
            ans[i] = str(base[2])
            base[2] += 3
        else:
            ans[i] = str(base[oe])
            base[oe] += 3
elif even_len <= N//3:
    for i, oe in enumerate(ans):
        if oe == 0:
            ans[i] = str(base[2])
            base[2] += 3
        elif base[0] <= N:
            ans[i] = str(base[0])
            base[0] += 3
        elif base[1] <= N:
            ans[i] = str(base[1])
            base[1] += 3
        else:
            ans[i] = str(base[2])
            base[2] += 3
elif odd_len <= N//3:
    for i, oe in enumerate(ans):
        if oe == 1:
            ans[i] = str(base[2])
            base[2] += 3
        elif base[0] <= N:
            ans[i] = str(base[0])
            base[0] += 3
        elif base[1] <= N:
            ans[i] = str(base[1])
            base[1] += 3
        else:
            ans[i] = str(base[2])
            base[2] += 3
print(' '.join(ans))

Buggy submission (verdict: Runtime Error):
from sys import setrecursionlimit

setrecursionlimit(10**6)

def dfs(i, depth):
    visited[i] = True
    oe = depth%2
    ans[i-1] = oe
    oe_index[oe] += [i]
    nodes = tree[i]
    for node in nodes:
        if visited[node]:
            continue
        dfs(node, depth+1)
    


N = int(input())
ans = [-1]*N
tree = [[] for i in range(N+1)]
visited = [False]*(N+1)
oe_index = [[], []]
base = [1, 2, 3]
for i in range(N-1):
    a, b = map(int, input().split())
    tree[a] += [b]
    tree[b] += [a]
dfs(1, 0)
even_len = len(oe_index[0])
odd_len = len(oe_index[1])
if even_len > N//3 and odd_len > N//3:
    for i, oe in enumerate(ans):
        if base[oe] > N:
            ans[i] = str(base[2])
            base[2] += 3
        else:
            ans[i] = str(base[oe])
            base[oe] += 3
elif even_len <= N//3:
    for i in oe_index[0]:
        ans[i] = str(base[2])
        base[2] += 3
    for i in oe_index[1]:
        if base[0] <= N:
            ans[i] = str(base[0])
            base[0] += 3
        else:
            ans[i] = str(base[1])
            base[1] += 3
elif odd_len <= N//3:
    for i in oe_index[1]:
        ans[i] = str(base[2])
        base[2] += 3
    for i in oe_index[0]:
        if base[0] <= N:
            ans[i] = str(base[0])
            base[0] += 3
        else:
            ans[i] = str(base[1])
            base[1] += 3
print(' '.join(ans))

Buggy submission (verdict: Runtime Error):
from sys import setrecursionlimit

setrecursionlimit(10**6)

def dfs(i, depth):
    visited[i] = True
    oe = depth%2
    ans[i-1] = oe
    count[oe] += 1
    nodes = tree[i]
    for node in nodes:
        if visited[node]:
            continue
        dfs(node, depth+1)
    


N = int(input())
ans = [-1]*N
tree = [[] for i in range(N+1)]
visited = [False]*(N+1)
count = [0, 0]
base = [1, 2, 3]
for i in range(N-1):
    a, b = map(int, input().split())
    tree[a] += [b]
    tree[b] += [a]
dfs(1, 0+1)
if count[0] >= count[1]:
    for i, oe in enumerate(ans):
        if base[oe] > N:
            ans[i-1] = str(base[2])
            base[2] += 3
        else:
            ans[i-1] = str(base[oe])
            base[oe] += 3
else:
    for i, oe in enumerate(ans):
        oe = (oe+1)%2
        if base[oe] > N:
            ans[i] = str(base[2])
            base[2] += 3
        else:
            ans[i] = str(base[oe])
            base[oe] += 3
print(' '.join(ans))