Problem p02749 (python)

Fixed submission:
import sys
input = sys.stdin.readline

n=int(input())
a= [list(map(int, input().split())) for i in range(n-1)]

b = [[] for i in range(n)]
a1=[]
a2=[]
ans=[0]*n
for i in range(n-1):
    b[a[i][0]-1].append(a[i][1]-1)
    b[a[i][1]-1].append(a[i][0]-1)

import collections
from collections import deque
def tree(s):

    INF=-10**9
    dis = [INF for i in range(n)]
    dis[s]=0
    def bfs():
        d = deque()
        d.append(s)
        a1.append(s)

        while len(d):
            x = d.popleft()

            for i in range(len(b[x])):
                y=b[x][i]
                if dis[y] == INF:
                    d.append(y)
                    dis[y]=dis[x]+1
                    if dis[y]%2==0:
                        a1.append(y)
                    else:
                        a2.append(y)
        return a1,a2
    return bfs()

tree(0)

j = 0
k = 0
if min(len(a1),len(a2))<=n//3:
    if len(a1)<=n//3:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a1)-1:
                ans[a1[j]]=i
                j+=1
            else:
                ans[a2[k]]=i
                k+=1
    else:
        for i in range(1,n+1):
            if i<=len(a2)*3 and i%3==0 and j<=len(a2)-1:
                ans[a2[j]]=i
                j+=1
            else:
                ans[a1[k]]=i
                k+=1

else:
    for i in range(1,n+1):
        if i%3==1:
            ans[a1[j]]=i
            j+=1
        elif i%3==2:
            ans[a2[k]] = i
            k += 1
        else:
            continue

    for i in range(3,n+1,3):
        if j<=len(a1)-1:
            ans[a1[j]]=i
            j+=1
        else:
            ans[a2[k]] = i
            k += 1


print(*ans)

Buggy submission (verdict: Runtime Error):
n=int(input())
a= [list(map(int, input().split())) for i in range(n-1)]

b = [[] for i in range(n)]
a1=[]
a2=[]
ans=[0]*n
for i in range(n-1):
    b[a[i][0]-1].append(a[i][1]-1)
    b[a[i][1]-1].append(a[i][0]-1)

import collections
from collections import deque
def tree(s):

    INF=-10**9
    dis = [INF for i in range(n)]
    dis[s]=0
    def bfs():
        d = deque()
        d.append(s)
        a1.append(s)

        while len(d):
            x = d.popleft()

            for i in range(len(b[x])):
                y=b[x][i]
                if dis[y] == INF:
                    d.append(y)
                    dis[y]=dis[x]+1
                    if dis[y]%2==0:
                        a1.append(y)
                    else:
                        a2.append(y)
        return a1,a2
    return bfs()

tree(0)

j = 0
k = 0
if min(len(a1),len(a2))<=n//3:
    if len(a1)<=n//3:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a1):
                ans[a1[j]]=i
                j+=1
            else:
                ans[a2[k]]=i
                k+=1
    else:
        for i in range(1,n+1):
            if i<=len(a2)*3 and i%3==0 and j<=len(a2):
                ans[a2[j]]=i
                j+=1
            else:
                ans[a1[k]]=i
                k+=1

else:
    for i in range(1,n+1):
        if (i%3==0 or i%3==1) and j<=len(a1):
            ans[a1[j]]=i
            j+=1
        else:
            ans[a2[k]] = i
            k += 1

print(*ans)

Buggy submission (verdict: Runtime Error):
n=int(input())
a= [list(map(int, input().split())) for i in range(n-1)]

b = [[] for i in range(n)]
a1=[]
a2=[]
ans=[0]*n
for i in range(n-1):
    b[a[i][0]-1].append(a[i][1]-1)
    b[a[i][1]-1].append(a[i][0]-1)

import collections
from collections import deque
def tree(s):

    INF=-10**9
    dis = [INF for i in range(n)]
    dis[s]=0
    def bfs():
        d = deque()
        d.append(s)
        a1.append(s)

        while len(d):
            x = d.popleft()

            for i in range(len(b[x])):
                y=b[x][i]
                if dis[y] == INF:
                    d.append(y)
                    dis[y]=dis[x]+1
                    if dis[y]%2==0:
                        a1.append(y)
                    else:
                        a2.append(y)
        return a1,a2
    return bfs()

tree(0)

j = 0
k = 0
if min(len(a1),len(a2))<=n//3:
    if len(a1)<=n//3:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a1):
                ans[a1[j]]=i
                j+=1
            else:
                ans[a2[k]]=i
                k+=1
    else:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a2):
                ans[a2[j]]=i
                j+=1
            else:
                ans[a1[k]]=i
                k+=1

else:
    for i in range(1,n+1):
        if (i%3==0 or i%3==1) and j<=len(a1):
            ans[a1[j]]=i
            j+=1
        else:
            ans[a2[k]] = i
            k += 1

print(*ans)

Buggy submission (verdict: Runtime Error):
n=int(input())
a= [list(map(int, input().split())) for i in range(n-1)]

b = [[] for i in range(n)]
a1=[]
a2=[]
ans=[0]*n
for i in range(n-1):
    b[a[i][0]-1].append(a[i][1]-1)
    b[a[i][1]-1].append(a[i][0]-1)

import collections
from collections import deque
def tree(s):

    INF=-10**9
    dis = [INF for i in range(n)]
    dis[s]=0
    def bfs():
        d = deque()
        d.append(s)
        a1.append(s)

        while len(d):
            x = d.popleft()

            for i in range(len(b[x])):
                y=b[x][i]
                if dis[y] == INF:
                    d.append(y)
                    dis[y]=dis[x]+1
                    if dis[y]%2==0:
                        a1.append(y)
                    else:
                        a2.append(y)
        return a1,a2
    return bfs()

tree(0)

j = 0
k = 0
if min(len(a1),len(a2))<=n//3:
    if a1<=n//3:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a1):
                ans[a1[j]]=i
                j+=1
            else:
                ans[a2[k]]=i
                k+=1
    else:
        for i in range(1,n+1):
            if i<=len(a1)*3 and i%3==0 and j<=len(a2):
                ans[a2[j]]=i
                j+=1
            else:
                ans[a1[k]]=i
                k+=1

else:
    for i in range(1,n+1):
        if (i%3==0 or i%3==1) and j<=len(a1):
            ans[a1[j]]=i
            j+=1
        else:
            ans[a2[k]] = i
            k += 1

print(*ans)