Problem p02749 (Python)

Buggy submission (status: Runtime Error):
n = int(input())
ab = [list(map(int, input().split())) for _ in range(n - 1)]

adj = [[] for _ in range(n)]
for a, b in ab:
    a -= 1
    b -= 1
    adj[a].append(b)
    adj[b].append(a)


class Tree:
    WHITE = 0
    GRAY = -1
    BLACK = 1

    def __init__(self, adj):
        n = len(adj)
        self.adj = adj
        self.colors = [self.WHITE] * n
        self.depths = [-1] * n
        self.depth = 0

    def init(self):
        self.__init__(self.adj)

    def dfs(self, u):
        if self.colors[u] == self.BLACK:
            return

        self.colors[u] = self.GRAY
        self.depths[u] = self.depth

        for v in self.adj[u]:
            if self.colors[v] == self.WHITE:
                self.depth += 1
                self.dfs(v)
                self.depth -= 1

        self.colors[u] = self.BLACK


tree = Tree(adj)
tree.dfs(0)

ev = [i for i, e in enumerate(tree.depths) if e % 2 == 0]
od = [i for i, e in enumerate(tree.depths) if e % 2]

l_ev = len(ev)
l_od = len(od)

mod0 = [e for e in range(1, n + 1) if e % 3 == 0]
mod1 = [e for e in range(1, n + 1) if e % 3 == 1]
mod2 = [e for e in range(1, n + 1) if e % 3 == 2]

ans = [0] * n


def f(li1, li2):
    while li1 and li2:
        i = li1.pop()
        num = li2.pop()
        ans[i] = num


if l_ev > n / 3 and l_od > n / 3:
    f(ev, mod1)
    f(od, mod2)
    f(ev, mod0)
    f(od, mod0)

elif l_ev <= n / 3:
    f(ev, mod0)
    f(od, mod0)
    f(od, mod1)
    f(od, mod2)

else:
    f(od, mod0)
    f(ev, mod0)
    f(ev, mod1)
    f(ev, mod2)

print(*ans)


Fixed submission:
import sys

sys.setrecursionlimit(10**6)

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

adj = [[] for _ in range(n)]
for a, b in ab:
    a -= 1
    b -= 1
    adj[a].append(b)
    adj[b].append(a)


class Tree:
    WHITE = 0
    GRAY = -1
    BLACK = 1

    def __init__(self, adj):
        n = len(adj)
        self.adj = adj
        self.colors = [self.WHITE] * n
        self.depths = [-1] * n
        self.depth = 0

    def init(self):
        self.__init__(self.adj)

    def dfs(self, u):
        if self.colors[u] == self.BLACK:
            return

        self.colors[u] = self.GRAY
        self.depths[u] = self.depth

        for v in self.adj[u]:
            if self.colors[v] == self.WHITE:
                self.depth += 1
                self.dfs(v)
                self.depth -= 1

        self.colors[u] = self.BLACK


tree = Tree(adj)
tree.dfs(0)

ev = [i for i, e in enumerate(tree.depths) if e % 2 == 0]
od = [i for i, e in enumerate(tree.depths) if e % 2]

l_ev = len(ev)
l_od = len(od)

mod0 = [e for e in range(1, n + 1) if e % 3 == 0]
mod1 = [e for e in range(1, n + 1) if e % 3 == 1]
mod2 = [e for e in range(1, n + 1) if e % 3 == 2]

ans = [0] * n


def f(li1, li2):
    while li1 and li2:
        i = li1.pop()
        num = li2.pop()
        ans[i] = num


if l_ev > n / 3 and l_od > n / 3:
    f(ev, mod1)
    f(od, mod2)
    f(ev, mod0)
    f(od, mod0)

elif l_ev <= n / 3:
    f(ev, mod0)
    f(od, mod0)
    f(od, mod1)
    f(od, mod2)

else:
    f(od, mod0)
    f(ev, mod0)
    f(ev, mod1)
    f(ev, mod2)

print(*ans)