Problem p02749 (python)

Fixed submission:
class Tree:
    def __init__(self, N):
        self.N = N # #vertices
        self.D = (self.N-1).bit_length() # #iterations in doubling
        self.E = [{} for _ in range(N)]
    
    def add_edge(self, init, end, weight, undirected=False):
        self.E[init][end] = weight
        if undirected: self.E[end][init] = weight
    
    def lca_initialize(self):
        self.dfs(0)
        self.doubling()
    
    def dfs(self, root):
        # parents[v]: the parent of the vertex v
        # depth[v]: the depth of the vertex v from the root
        # dist[v]: the distance of the vertex v from the root.
        self.parent = [None] * self.N
        self.depth = [None] * self.N
        self.dist = [float('inf')] * self.N
        self.parent[root] = -1; self.depth[root] = 0; self.dist[root] = 0;
        stack = [root]
        while stack:
            v = stack.pop()
            for u, d in self.E[v].items():
                if self.depth[u] is not None: continue
                self.parent[u] = v
                self.depth[u] = self.depth[v] + 1
                self.dist[u] = self.dist[v] + d
                stack.append(u)
    
    def doubling(self):
        # O(N log N) time
        self.next = [[-1] * self.N for _ in range(self.D)]
        for v in range(self.N):
            self.next[0][v] = self.parent[v]
        for d in range(self.D - 1):
            for v in range(self.N):
                if self.next[d][v] != -1:
                    self.next[d + 1][v] = self.next[d][self.next[d][v]]
    
    def lca(self, u, v):
        # O(log N) time
        # the depth of v is set to be no less than that of u
        if self.depth[u] > self.depth[v]: u, v = v, u
        # find the ancestor of v with the same depth as that of u
        diff = self.depth[v] - self.depth[u]
        for i in range(diff.bit_length()):
            if diff & (1 << i): v = self.next[i][v]
        if u == v: return u
        for i in range(self.depth[u].bit_length() - 1, -1, -1):
            pu, pv = self.next[i][u], self.next[i][v]
            if pu != pv: u, v = pu, pv
        return self.next[0][u]

N = int(input())
tr = Tree(N)
for _ in range(N-1):
    a, b = map(int, input().split())
    tr.add_edge(a-1, b-1, 1, undirected=True)
tr.dfs(0)
d_odd, d_even = [], []
rem = [[] for _ in range(3)]
for i in range(N):
    if tr.depth[i] % 2 == 0: d_even.append(i)
    else: d_odd.append(i)
    rem[(i+1)%3].append(i+1)
P = [0] * N
if len(d_odd) < len(d_even):
    if len(d_odd) <= len(rem[0]):
        for i in d_odd: P[i] = rem[0].pop()
    else:
        for i in d_odd:
            if rem[2]:
                P[i] = rem[2].pop()
            else:
                P[i] = rem[0].pop()
    temp = rem[2] + rem[0] + rem[1]
    for i in d_even: P[i] = temp.pop()
else:
    if len(d_even) <= len(rem[0]):
        for i in d_even: P[i] = rem[0].pop()
    else:
        for i in d_even:
            if rem[2]:
                P[i] = rem[2].pop()
            else:
                P[i] = rem[0].pop()
    temp = rem[2] + rem[0] + rem[1]
    for i in d_odd: P[i] = temp.pop()
print(*P)

Buggy submission (verdict: Runtime Error):
class Tree:
    def __init__(self, N):
        self.N = N # #vertices
        self.D = (self.N-1).bit_length() # #iterations in doubling
        self.E = [{} for _ in range(N)]
    
    def add_edge(self, init, end, weight, undirected=False):
        self.E[init][end] = weight
        if undirected: self.E[end][init] = weight
    
    def lca_initialize(self):
        self.dfs(0)
        self.doubling()
    
    def dfs(self, root):
        # parents[v]: the parent of the vertex v
        # depth[v]: the depth of the vertex v from the root
        # dist[v]: the distance of the vertex v from the root.
        self.parent = [None] * self.N
        self.depth = [None] * self.N
        self.dist = [float('inf')] * self.N
        self.parent[root] = -1; self.depth[root] = 0; self.dist[root] = 0;
        stack = [root]
        while stack:
            v = stack.pop()
            for u, d in self.E[v].items():
                if self.depth[u] is not None: continue
                self.parent[u] = v
                self.depth[u] = self.depth[v] + 1
                self.dist[u] = self.dist[v] + d
                stack.append(u)
    
    def doubling(self):
        # O(N log N) time
        self.next = [[-1] * self.N for _ in range(self.D)]
        for v in range(self.N):
            self.next[0][v] = self.parent[v]
        for d in range(self.D - 1):
            for v in range(self.N):
                if self.next[d][v] != -1:
                    self.next[d + 1][v] = self.next[d][self.next[d][v]]
    
    def lca(self, u, v):
        # O(log N) time
        # the depth of v is set to be no less than that of u
        if self.depth[u] > self.depth[v]: u, v = v, u
        # find the ancestor of v with the same depth as that of u
        diff = self.depth[v] - self.depth[u]
        for i in range(diff.bit_length()):
            if diff & (1 << i): v = self.next[i][v]
        if u == v: return u
        for i in range(self.depth[u].bit_length() - 1, -1, -1):
            pu, pv = self.next[i][u], self.next[i][v]
            if pu != pv: u, v = pu, pv
        return self.next[0][u]

N = int(input())
tr = Tree(N)
for _ in range(N-1):
    a, b = map(int, input().split())
    tr.add_edge(a-1, b-1, 1, undirected=True)
tr.dfs(0)
d_odd, d_even = [], []
rem = [[] for _ in range(3)]
for i in range(N):
    if tr.depth[i] % 2 == 0: d_even.append(i)
    else: d_odd.append(i)
    rem[(i+1)%3].append(i+1)
P = [0] * N
flag = True
if len(d_odd) < len(d_even):
    for i in d_odd:
        if rem[2]:
            P[i] = rem[2].pop()
        elif rem[0]:
            P[i] = rem[0].pop()
        else:
            flag = False; break
    temp = rem[0] + rem[1]
    if flag:
        for i in d_even:
            if temp:
                P[i] = temp.pop()
            else:
                flag = False; break
else:
    for i in d_even:
        if rem[2]:
            P[i] = rem[2].pop()
        elif rem[0]:
            P[i] = rem[0].pop()
        else:
            flag = False; break
    temp = rem[0] + rem[1]
    if flag:
        for i in d_odd:
            if temp:
                P[i] = temp.pop()
            else:
                flag = False; break
print(*P if flag else -1)

Buggy submission (verdict: Runtime Error):
class Tree:
    def __init__(self, N):
        self.N = N # #vertices
        self.D = (self.N-1).bit_length() # #iterations in doubling
        self.E = [{} for _ in range(N)]
    
    def add_edge(self, init, end, weight, undirected=False):
        self.E[init][end] = weight
        if undirected: self.E[end][init] = weight
    
    def lca_initialize(self):
        self.dfs(0)
        self.doubling()
    
    def dfs(self, root):
        # parents[v]: the parent of the vertex v
        # depth[v]: the depth of the vertex v from the root
        # dist[v]: the distance of the vertex v from the root.
        self.parent = [None] * self.N
        self.depth = [None] * self.N
        self.dist = [float('inf')] * self.N
        self.parent[root] = -1; self.depth[root] = 0; self.dist[root] = 0;
        stack = [root]
        while stack:
            v = stack.pop()
            for u, d in self.E[v].items():
                if self.depth[u] is not None: continue
                self.parent[u] = v
                self.depth[u] = self.depth[v] + 1
                self.dist[u] = self.dist[v] + d
                stack.append(u)
    
    def doubling(self):
        # O(N log N) time
        self.next = [[-1] * self.N for _ in range(self.D)]
        for v in range(self.N):
            self.next[0][v] = self.parent[v]
        for d in range(self.D - 1):
            for v in range(self.N):
                if self.next[d][v] != -1:
                    self.next[d + 1][v] = self.next[d][self.next[d][v]]
    
    def lca(self, u, v):
        # O(log N) time
        # the depth of v is set to be no less than that of u
        if self.depth[u] > self.depth[v]: u, v = v, u
        # find the ancestor of v with the same depth as that of u
        diff = self.depth[v] - self.depth[u]
        for i in range(diff.bit_length()):
            if diff & (1 << i): v = self.next[i][v]
        if u == v: return u
        for i in range(self.depth[u].bit_length() - 1, -1, -1):
            pu, pv = self.next[i][u], self.next[i][v]
            if pu != pv: u, v = pu, pv
        return self.next[0][u]

N = int(input())
tr = Tree(N)
for _ in range(N-1):
    a, b = map(int, input().split())
    tr.add_edge(a-1, b-1, 1, undirected=True)
tr.dfs(0)
d_odd, d_even = [], []
rem = [[] for _ in range(3)]
for i in range(N):
    if tr.depth[i] % 2 == 0: d_even.append(i)
    else: d_odd.append(i)
    rem[(i+1)%3].append(i+1)
P = [0] * N
if len(d_odd) < len(d_even):
    for i in d_odd:
        if rem[2]:
            P[i] = rem[2].pop()
        elif rem[0]:
            P[i] = rem[0].pop()
        else:
            P[i] = rem[1].pop()
    temp = rem[0] + rem[1]
    for i in d_even: P[i] = temp.pop()
else:
    for i in d_even:
        if rem[2]:
            P[i] = rem[2].pop()
        elif rem[0]:
            P[i] = rem[0].pop()
        else:
            P[i] = rem[1].pop()
    temp = rem[0] + rem[1]
    for i in d_odd: P[i] = temp.pop()
print(*P)