Problem p02749 (python)

Fixed submission:
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    vs = [[] for _ in range(2)]
    vs[0].append(0)

    while dq:
        v = dq.popleft()
        p = parity[v]
        np = 1 ^ p
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = np
            vs[np].append(u)
            dq.append(u)

    # SOLVE
    k = N // 3  # 3の倍数の個数

    g1 = iter(range(1, N + 1, 3))
    g2 = iter(range(2, N + 1, 3))
    g3 = iter(range(3, N + 1, 3))
    if len(vs[0]) <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        for v in vs[0]:
            parity[v] = next(g3)  # g3の方がvs[0]より数多い

        g123 = chain(g1, g2, g3)

        for v in vs[1]:
            parity[v] = next(g123)

    elif len(vs[1]) <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        for v in vs[1]:
            parity[v] = next(g3)  # g3の方がvs[1]より数多い

        g123 = chain(g1, g2, g3)

        for v in vs[0]:
            parity[v] = next(g123)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める

        g13 = chain(g1, g3)
        for v in vs[0]:
            parity[v] = next(g13)
            # k<=len(g1)<=k+1,len(vs[0])>=k+1
            # 右辺がg1だけだと不足する恐れ

        g23 = chain(g2, g3)
        for v in vs[1]:
            parity[v] = next(g23)
            # k<=len(g2)<=k+1,len(vs[1])>=k+1
            # 右辺がg2だけだと不足する恐れ

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    bipartite = [-1] * N
    bipartite[0] = 0

    dq = deque([0])
    while dq:
        v = dq.popleft()
        nb = bipartite[v] ^ 1
        for u in graph[v]:
            if ~bipartite[u]: continue
            bipartite[u] = nb
            dq.append(u)

    s0 = bipartite.count(0)
    s1 = N - s0

    it1 = iter(range(1, N + 1, 3))
    it2 = iter(range(2, N + 1, 3))
    it3 = iter(range(3, N + 1, 3))
    it12 = chain(it1, it2)
    it13 = chain(it1, it3)
    it23 = chain(it2, it3)

    ans = [-1] * N
    K = N // 3
    if s0 <= K:
        for v in range(N):
            if bipartite[v] == 0:
                ans[v] = next(it3)
            else:
                ans[v] = next(it12)
    elif s1 <= K:
        for v in range(N):
            if bipartite[v] == 1:
                ans[v] = next(it3)
            else:
                ans[v] = next(it12)
    else:
        for v in range(N):
            if bipartite[v] == 0:
                ans[v] = next(it13)
            else:
                ans[v] = next(it23)

    print(*ans)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g3r = range(N - N % 3, 3 - 1, -3)
        g12r3 = chain(g1, g2, g3r)

        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12r3)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g3r = range(N - N % 3, 3 - 1, -3)
        g12r3 = chain(g1, g2, g3r)

        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12r3)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g3r = range(N - N % 3, 3 - 1, -3)
        g13 = chain(g1, g3)
        g2r3 = chain(g2, g3r)

        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g12 = chain(g1, g2)

        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g12 = chain(g1, g2)

        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        g1 = range(1, N + 1, 3)
        g2 = range(2, N + 1, 3)
        g3 = range(3, N + 1, 3)
        g3r = range(N - N % 3, 3 - 1, -3)
        g13 = chain(g1, g3)
        g2r3 = chain(g2, g3r)

        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    g1 = range(1, N + 1, 3)
    g2 = range(2, N + 1, 3)
    g3 = range(3, N + 1, 3)
    g3r = range((N // 3) * 3, 3 - 1, -3)
    g12 = chain(g1, g2)
    g13 = chain(g1, g3)
    g2r3 = chain(g2, g3r)
    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    g1 = range(1, N + 1, 3)
    g2 = range(2, N + 1, 3)
    g3 = range(3, N + 1, 3)
    g3r = range(N, 3 - 1, -3)
    g12 = chain(g1, g2)
    g13 = chain(g1, g3)
    g2r3 = chain(g2, g3r)
    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain, tee
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    g1 = range(1, N + 1, 3)
    g2 = range(2, N + 1, 3)
    g3 = range(3, N + 1, 3)
    g3f, g3r = tee(g3, 2)
    g12 = chain(g1, g2)
    g13 = chain(g1, g3f)
    g2r3 = chain(g2, reversed(g3r))
    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# 完全二部グラフで大丈夫なら全部大丈夫(最悪で大丈夫かどうか)
# 簡単のためN=3kとする(3の倍数以外でもほとんど同じなので)
# 0,1,2がk個ずつ

# 1-1,2-2間に辺があるとダメ(これだけを避ければよい)
#   1-2,1-0,2-1,2-0,0-0はOKなので

# 二部グラフで片側の頂点が2k以上(もう一方がk個以下)の場合
#   2k以上の方に1,2を割り当てて、反対側をすべて0にする

# 両方とも(k,2k)の場合
# 両方k以上なのでLを1、Rを2で埋めて残りを0で埋める


def main():
    from collections import deque
    from itertools import chain
    import sys
    input = sys.stdin.readline

    # INPUT
    N = int(input())

    graph = tuple(set() for _ in range(N))
    for _ in range(N - 1):
        a, b = (int(x) - 1 for x in input().split())
        graph[a].add(b)
        graph[b].add(a)

    # BFS
    dq = deque([0])
    parity = [-1] * N
    parity[0] = 0

    while dq:
        v = dq.popleft()
        p = parity[v]
        for u in graph[v]:
            if ~parity[u]: continue
            parity[u] = 1 ^ p
            dq.append(u)

    # SOLVE
    count_0 = parity.count(0)
    k = N // 3  # 3の倍数の個数

    g1 = range(1, N + 1, 3)
    g2 = range(2, N + 1, 3)
    g3 = range(3, N + 1, 3)
    g12 = chain(g1, g2)
    g13 = chain(g1, g3)
    g2r3 = chain(g2, reversed(g3))
    if count_0 <= k:
        # parity=0を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 0 else next(g12)

    elif N - count_0 <= k:
        # parity=1を3の倍数で埋める
        # 残りを余った数で埋める
        for v in range(N):
            parity[v] = next(g3) if parity[v] == 1 else next(g12)

    else:
        # parity=0を(mod3=1)で埋める
        # parity=1を(mod3=2)で埋める
        # 残りを3の倍数で埋める
        for v in range(N):
            if parity[v] == 0:
                parity[v] = next(g13)
            else:
                parity[v] = next(g2r3)

    print(*parity)


if __name__ == '__main__':
    main()