Problem p02749 (python)

Fixed submission:
import sys
input = sys.stdin.readline
n = int(input())
tree = [[] for i in range(n+1)]
for i in range(n-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)


ans = [-1]*(n+1)
q = [1]
ans[1] = 0
while q:
    now = q.pop()
    for next in tree[now]:
        if ans[next] == -1:
            q.append(next)
            ans[next] = (ans[now]+1)%2

if ans.count(1) <= n//3:
    three = set([3*i+3 for i in range(n//3)])
    onetwo = set([i+1 for i in range(n)])
    onetwo = onetwo - three
    three = list(three)
    three.sort()
    onetwo = list(onetwo)
    onetwo.sort()
    three_ind = 0
    onetwo_ind = 0
    for i in range(1, n+1):
        if ans[i] == 1:
            ans[i] = three[three_ind]
            three_ind += 1
        else:
            if onetwo_ind >= len(onetwo):
                ans[i] = three[three_ind]
                three_ind += 1
            else:
                ans[i] = onetwo[onetwo_ind]
                onetwo_ind += 1

elif ans.count(0) <= n//3:
    three = set([3*i+3 for i in range(n//3)])
    onetwo = set([i+1 for i in range(n)])
    onetwo = onetwo - three
    three = list(three)
    three.sort()
    onetwo = list(onetwo)
    onetwo.sort()
    three_ind = 0
    onetwo_ind = 0
    for i in range(1, n+1):
        if ans[i] == 0:
            ans[i] = three[three_ind]
            three_ind += 1
        else:
            if onetwo_ind >= len(onetwo):
                ans[i] = three[three_ind]
                three_ind += 1
            else:
                ans[i] = onetwo[onetwo_ind]
                onetwo_ind += 1

else:
    zero = 0
    one = 0
    two = 0
    for i in range(1, n+1):
        if ans[i] == 0:
            ans[i] = zero*3+1
            zero += 1
        elif ans[i] == 1:
            ans[i] = one*3+2
            one += 1

    three = 0
    for i in range(1, n+1):
        if ans[i] > n:
            ans[i] = three*3+3
            three += 1

print(" ".join(map(str, ans[1:])))

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
n = int(input())
tree = [[] for i in range(n+1)]
for i in range(n-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)


ans = [-1]*(n+1)
q = [1]
ans[1] = 0
while q:
    now = q.pop()
    for next in tree[now]:
        if ans[next] == -1:
            q.append(next)
            ans[next] = (ans[now]+1)%2

if ans.count(1) <= n//3:
    three = set([3*i+3 for i in range(n//3)])
    onetwo = set([i+1 for i in range(n)])
    onetwo = onetwo - three
    three = list(three)
    three.sort()
    onetwo = list(onetwo)
    onetwo.sort()
    three_ind = 0
    onetwo_ind = 0
    for i in range(1, n+1):
        if ans[i] == 1:
            ans[i] = three[three_ind]
            three_ind += 1
        else:
            ans[i] = onetwo[onetwo_ind]
            onetwo_ind += 1

elif ans.count(0) <= n//3:
    three = set([3*i+3 for i in range(n//3)])
    onetwo = set([i+1 for i in range(n)])
    onetwo = onetwo - three
    three = list(three)
    three.sort()
    onetwo = list(onetwo)
    onetwo.sort()
    three_ind = 0
    onetwo_ind = 0
    for i in range(1, n+1):
        if ans[i] == 0:
            ans[i] = three[three_ind]
            three_ind += 1
        else:
            ans[i] = onetwo[onetwo_ind]
            onetwo_ind += 1

else:
    zero = 0
    one = 0
    two = 0
    for i in range(1, n+1):
        if ans[i] == 0:
            ans[i] = zero*3+1
            zero += 1
        elif ans[i] == 1:
            ans[i] = one*3+2
            one += 1

    three = 0
    for i in range(1, n+1):
        if ans[i] > n:
            ans[i] = three*3+3
            three += 1

print(" ".join(map(str, ans[1:])))