Problem p02749 (python)

Fixed submission:
N = int(input())
data = [[] for i in range(N)]

for i in range(N-1):
    a,b = map(int,input().split())
    a -= 1
    b -= 1
    data[a].append(b)
    data[b].append(a)

stack = [0]
d = [-1] * N
d[0] = 0

while stack:
    x = stack.pop()
    for y in data[x]:
        if d[y] != -1:
            continue
        d[y] = d[x] + 1
        stack.append(y)

mod = [[],[]]
for i in range(N):
    mod[d[i]%2].append(i)

ans = [0]*N

p = 1
q = 2
r = 3
if len(mod[0]) > N//3 and len(mod[1]) > N//3:
    for k in mod[0]:
        if p <= N:
            ans[k] = p
            p += 3
        else:
            ans[k] = r
            r += 3
    for l in mod[1]:
        if q <= N:
            ans[l] = q
            q += 3
        else:
            ans[l] = r
            r += 3
    
elif len(mod[0]) > N//3:
    for k in mod[0]:
        if p <= N:
            ans[k] = p
            p += 3
        elif q <= N:
            ans[k] = q
            q += 3
        else:
            ans[k] = r
            r += 3
    for l in mod[1]:
        ans[l] = r
        r += 3

else:
    for k in mod[1]:
        if p <= N:
            ans[k] = p
            p += 3
        elif q <= N:
            ans[k] = q
            q += 3
        else:
            ans[k] = r
            r += 3
    for l in mod[0]:
        ans[l] = r
        r += 3

print(" ".join(map(str,ans)))

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

N = int(input())
data1 = []
data2 = [[] for i in range(N)]
data3 = [set() for i in range(N)]

for i in range(N-1):
    a,b = map(int,input().split())
    a -= 1
    b -= 1
    data1.append([a,b])
    data2[a].append(b)
    data2[b].append(a)

for i in range(N-1):
    a,b = data1[i]
    for x in data2[a]:
        for y in data2[b]:
            if y != a and x != b:
                data3[x].add(y)
                data3[y].add(x)

for i in range(N):
    data3[i] = list(data3[i])

#####print(data3)

s = [0,0,0]
while len(data2[s[0]]) <= 1:
    s[0] += 1

s[1] = data2[s[0]][0]
s[2] = data2[s[0]][1]

stack = deque([[s[0],1],[s[1],1],[s[2],2]])
ans = [0] * N
ans[s[0]] = 1
ans[s[1]] = 4
ans[s[2]] = 2
count = [3,7,5]

det = 0
while stack:
    p = stack.pop()
    for q in data3[p[0]]:
        if ans[q] != 0:
            continue
        if count[2] <= N and p[1]%3 != 2:
            ans[q] = count[2]
            count[2] += 3
            stack.append([q,2])
        elif count[1] <= N and p[1]%3 != 1:
            ans[q] = count[1]
            count[1] += 3
            stack.append([q,1])
        else:
            if count[0] > N:
                det = -1
            ans[q] = count[0]
            count[0] += 3
            stack.append([q,0])
        #####print(ans,stack)

if det == -1:
    print(-1)
else:
    print(" ".join(map(str,ans)))

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

N = int(input())
data1 = []
data2 = [[] for i in range(N)]
data3 = [set() for i in range(N)]

for i in range(N-1):
    a,b = map(int,input().split())
    a -= 1
    b -= 1
    data1.append([a,b])
    data2[a].append(b)
    data2[b].append(a)

for i in range(N-1):
    a,b = data1[i]
    for x in data2[a]:
        for y in data2[b]:
            if y != a and x != b:
                data3[x].add(y)
                data3[y].add(x)

for i in range(N):
    data3[i] = list(data3[i])

#####print(data3)

s = [0,0,0]
while len(data2[s[0]]) <= 1:
    s[0] += 1

s[1] = data2[s[0]][0]
s[2] = data2[s[0]][1]

stack = deque([[s[0],1],[s[1],1],[s[2],2]])
ans = [0] * N
ans[s[0]] = 1
ans[s[1]] = 4
ans[s[2]] = 2
count = [3,7,5]

det = 0
while stack:
    p = stack.pop()
    for q in data3[p[0]]:
        if ans[q] != 0:
            continue
        if count[2] <= N and p[1]%3 != 2:
            ans[q] = count[2]
            count[2] += 3
            stack.append([q,2])
        elif count[1] <= N and p[1]%3 != 1:
            ans[q] = count[1]
            count[1] += 3
            stack.append([q,1])
        else:
            if count[0] > N:
                det = -1
            ans[q] = count[0]
            count[0] += 3
            stack.append([q,0])
        #####print(ans,stack)

if det == -1:
    print(-1)
else:
    print(" ".join(map(str,ans)))

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

N = int(input())
data1 = []
data2 = [[] for i in range(N)]
data3 = [set() for i in range(N)]

for i in range(N-1):
    a,b = map(int,input().split())
    a -= 1
    b -= 1
    data1.append([a,b])
    data2[a].append(b)
    data2[b].append(a)

for i in range(N-1):
    a,b = data1[i]
    for x in data2[a]:
        for y in data2[b]:
            if y != a and x != b:
                data3[x].add(y)
                data3[y].add(x)

for i in range(N):
    data3[i] = list(data3[i])

#####print(data3)

s = [0,0,0]
while len(data2[s[0]]) <= 1:
    s[0] += 1

s[1] = data2[s[0]][0]
s[2] = data2[s[0]][1]

stack = deque([[s[0],1],[s[1],1],[s[2],2]])
ans = [0] * N
ans[s[0]] = 1
ans[s[1]] = 4
ans[s[2]] = 2
count = [3,7,5]

while stack:
    p = stack.pop()
    for q in data3[p[0]]:
        if ans[q] != 0:
            continue
        if count[2] <= N and p[1]%3 != 2:
            ans[q] = count[2]
            count[2] += 3
            stack.append([q,2])
        elif count[1] <= N and p[1]%3 != 1:
            ans[q] = count[1]
            count[1] += 3
            stack.append([q,1])
        else:
            ans[q] = count[0]
            count[0] += 3
            stack.append([q,0])
        #####print(ans,stack)

print(" ".join(map(str,ans)))