Problem p02749 (python)

Fixed submission:
# coding: utf-8

N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

def run(cur, tr, dist, d=1, pre=-1, recur=1, save_recur = []):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        if recur >500:
            dic_ = {'cur':node, 'tr': tr, 'dist': dist,
                              'd': (d + 1) % 2, 'pre': cur, 'recur': recur + 1}
            save_recur.append(dic_)
            return None
        run(node, tr, dist, (d+1)%2, pre=cur, recur=recur+1, save_recur=save_recur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    save_rec = []
    run(1, tr, dist, save_recur=save_rec)

    while True:
        save_rec2 = save_rec
        save_rec = []
        for dict in save_rec2:
            run(dict['cur'], dict['tr'], dict['dist'], dict['d'],
                dict['pre'], dict['recur'],save_recur = save_rec)
        if save_rec == []:break


    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) <= N/3) | (len(dist[0]) >= 2*N/3):
        bool = int(len(dist[0]) <= N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem1_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        bool = int((len(dist[0]) > len(dist[1])))
        for alpha in range(n):
            even = dist[bool][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1-bool][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]

        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8

N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

import numpy as np
a_ = np.arange(1,10000).reshape(-1,1)
b_ = np.arange(2,10001).reshape(-1,1)
ab = np.concatenate([a_, b_], axis=1)
N = 10000

def run(cur, tr, dist, d=1, pre=-1, recur=1, save_recur = []):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        if recur >500:
            dic_ = {'cur':node, 'tr': tr, 'dist': dist,
                              'd': (d + 1) % 2, 'pre': cur, 'recur': recur + 1}
            save_recur.append(dic_)
            return None
        run(node, tr, dist, (d+1)%2, pre=cur, recur=recur+1, save_recur=save_recur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    save_rec = []
    run(1, tr, dist, save_recur=save_rec)

    while True:
        save_rec2 = save_rec
        save_rec = []
        for dict in save_rec2:
            run(**dict, save_recur = save_rec)
        if save_rec == []:break


    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) <= N/3) | (len(dist[0]) >= 2*N/3):
        bool = int(len(dist[0]) <= N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem1_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        bool = int((len(dist[0]) > len(dist[1])))
        for alpha in range(n):
            even = dist[bool][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1-bool][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]

        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]


def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) <= N/3) | (len(dist[0]) >= 2*N/3):
        bool = int(len(dist[0]) <= N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem1_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        bool = int((len(dist[0]) > len(dist[1])))
        for alpha in range(n):
            even = dist[bool][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1-bool][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]

        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]


def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) <= N/3) | (len(dist[0]) >= 2*N/3):
        bool = int(len(dist[0]) <= N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem0_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        bool = int((len(dist[0]) > len(dist[1])))
        for alpha in range(n):
            even = dist[bool][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1-bool][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]

        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    a = [1,2,3,4,5]
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])
    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) <= N/3) | (len(dist[0]) >= 2*N/3):
        bool = int(len(dist[0]) < N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem0_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        for alpha in range(n):
            even = dist[0][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        bool = (len(dist[0]) > len(dist[1]))
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]
        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    a = [1,2,3,4,5]
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])
    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    rem1_max, rem2_max, rem0_max = int(N/3),int(N/3),int(N/3)
    if N % 3 == 1:
        rem1_max +=1
    if N % 3 == 2:
        rem1_max += 1
        rem2_max += 1
    if (len(dist[0]) < N/3) | (len(dist[0]) > 2*N/3):
        bool = int(len(dist[0]) < N/3)
        h0 = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3*i for i in range(h0, rem0_max+1)] + [3*i+1 for i in range(0, rem0_max)] + [3*i+2 for i in range(0, rem2_max)]
        for k, val in zip(dist[bool], nums):
            perm[k] = val
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        h1 = 0
        h2 = 0
        idx = 0
        for alpha in range(n):
            even = dist[0][idx]
            perm[even] = alpha * 3 + 1
            odd = dist[1][idx]
            perm[odd] = alpha * 3 + 2
            idx += 1
        bool = (len(dist[0]) > len(dist[1]))
        h0 = 1
        for k in dist[1-bool][idx:]:
            perm[k] = 3 * h0
            h0 += 1
        nums = [3 * i for i in range(h0, rem0_max+1)] + [3 * i + 1 for i in range(n, rem1_max)] + [3 * i + 2 for i in range(n, rem2_max)]
        for h, k in enumerate(dist[bool][idx:]):
            perm[k] = nums[h]
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
import numpy as np
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    nums = list(range(N+1))
    if (len(dist[0]) < N/3) | (len(dist[0]) > 2*N/3):
        bool = int(len(dist[0]) < N/3)
        h = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h
            h += 1
            nums.remove(3*h)
        h = 1
        for k in dist[bool]:
            perm[k] = nums[h]
            h += 1
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        for alpha in range(n):
            even = dist[0][0]
            dist[0].remove(even)
            perm[even] = alpha * 3 + 1
            nums.remove(alpha * 3 + 1)
            odd = dist[1][0]
            dist[1].remove(odd)
            perm[odd] = alpha * 3 + 2
            nums.remove(alpha * 3 + 2)
        bool = (len(dist[0]) > len(dist[1]))
        h = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h
            nums.remove(3 * h)
            h += 1
        h=1
        for k in dist[bool]:
            perm[k] = nums[h]
            h += 1
    print(" ".join(list(map(str, perm[1:]))))

Buggy submission (verdict: Runtime Error):
# coding: utf-8
import numpy as np
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]

def run(cur, tr, dist, d=1, pre=-1):
    for node in tr[cur]:
        if node == pre:continue
        dist[d].append(node)
        run(node, tr, dist, (d+1)%2, pre=cur)
    return None

if __name__ == '__main__':
    # generate dict
    tr = {}
    for line in ab:
        if line[0] not in tr.keys():
            tr[line[0]] = [line[1]]
        else:
            tr[line[0]].append(line[1])
        if line[1] not in tr.keys():
            tr[line[1]] = [line[0]]
        else:
            tr[line[1]].append(line[0])

    dist = {0: [1], 1: []}
    run(1, tr, dist)

    perm = [0] * (N+1)
    nums = list(range(N+1))
    if (len(dist[0]) < N/3) | (len(dist[0]) > 2*N/3):
        bool = int(len(dist[0]) < N/3)
        h = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h
            h += 1
            nums.remove(3*h)
        h = 1
        for k in dist[bool]:
            perm[k] = nums[h]
            h += 1
    else:
        n = int(N/3)
        if N % 3 == 2:
            n += 1
        for alpha in range(n):
            even = dist[0][0]
            dist[0].remove(even)
            perm[even] = alpha * 3 + 1
            nums.remove(alpha * 3 + 1)
            odd = dist[1][0]
            dist[1].remove(odd)
            perm[odd] = alpha * 3 + 2
            nums.remove(alpha * 3 + 2)
        bool = (len(dist[0]) > len(dist[1]))
        h = 1
        for k in dist[1-bool]:
            perm[k] = 3 * h
            h += 1
            nums.remove(3 * h)
        h=1
        for k in dist[bool]:
            perm[k] = nums[h]
            h += 1
    print(" ".join(list(map(str, perm[1:]))))