Problem p02750 (python)

Fixed submission:
from operator import itemgetter
import bisect


n, t = map(int, input().split())
info = [list(map(int, input().split())) for i in range(n)]
INF = t + 1
LIMIT = 60

# mrkt_a0: a == 0 の店
# mrkt_others: a != 0 の店
mrkt_a0 = []
mrkt_others = []

for i in range(n):
    a, b = info[i]
    if a == 0:
        mrkt_a0.append((a, b))
    else:
        comp = a / (1 + b)
        mrkt_others.append((a, b, comp))

mrkt_a0.sort(key=itemgetter(1))
mrkt_others.sort(key=itemgetter(2), reverse=True)

# mrkt_a0について
# もしmarkt_a0を尋ねるなら尋ねるのを後回しにしたほうが良い
times_a0 = [0]
for a, b in mrkt_a0:
    times_a0.append(times_a0[-1] + 1 + b)

# mrkt_othersについて
len_others = len(mrkt_others)
dp = [[INF] * LIMIT for i in range(len_others + 1)]
dp[0][0] = 0

for i in range(len_others):
    a, b, _ = mrkt_others[i]
    for cnt in range(LIMIT):
        # i番目の店を尋ねないとき
        dp[i + 1][cnt] = min(dp[i][cnt], dp[i + 1][cnt])
        # i番目の店を尋ねるとき
        if cnt + 1 < LIMIT:
            dp[i + 1][cnt + 1] = min(a * (dp[i][cnt] + 1) + b + dp[i][cnt] + 1, dp[i + 1][cnt + 1])

ans = 0
for cnt in range(LIMIT):
    if dp[-1][cnt] == t + 1:
        continue
    nokori = t - dp[-1][cnt]
    cnt_a0 = bisect.bisect_right(times_a0, nokori) - 1
    ans = max(ans, cnt + cnt_a0)

print(ans)

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


n = int(input())

info = [list(map(int, input().split())) for i in range(n - 1)]
tree = [[] for i in range(n)]
for i in range(n - 1):
    a, b = info[i]
    a -= 1
    b -= 1
    tree[a].append(b)
    tree[b].append(a)


root = 0
visited = [-1] * n
def dfs(root):
    q = deque([root])
    visited[root] = 0
    while q:
        v = q.pop()
        for nxt_v in tree[v]:
            if visited[nxt_v] != -1:
                continue
            else:
                visited[nxt_v] = (visited[v] + 1) % 2
                q.append(nxt_v)
dfs(0)
cnt0 = 0
cnt1 = 0
for i in range(n):
    if visited[i] == 0:
        cnt0 += 1
    else:
        cnt1 += 1

num1 = n // 3 
num2 = n // 3
num3 = n // 3
if n % 3 >= 1:
   num1 += 1
if n % 3 >= 2:
   num2 += 1
    
ans = [0] * n

set1 = set()
set2 = set()
set3 = set()
for i in range(1, n + 1):
    if i % 3 == 0:
        set3.add(i)
    elif i % 3 == 1:
        set1.add(i)
    elif i % 3 == 2:
        set2.add(i)
        
if cnt0 <= num3:
    for i in range(n):
        if visited[i] == 0:
            ans[i] = pop.set3()
        else:
            if set1:
                ans[i] = pop.set1() 
            elif set2:
                ans[i] = pop.set2()         
            elif set3:
                ans[i] = pop.set3()  

elif cnt1 <= num3:
    for i in range(n):
        if visited[i] == 1:
            ans[i] = pop.set3()
        else:
            if set1:
                ans[i] = pop.set1() 
            elif set2:
                ans[i] = pop.set2()         
            elif set3:
                ans[i] = pop.set3()  
else:
    num147 = 0
    num258 = 0
    num369 = 0
    for i in range(n):
        if visited[i] == 0:
            num = num147 * 3 + 1
            if num > n:
                num = num369 * 3 + 3
                ans[i] = num
                if num > n:
                    print(-1)
                    exit()
                set3.remove(num)
                num369 += 1
                continue
            ans[i] = num # 1, 4, 7, 10, ...
            set1.remove(num)
            num147 += 1
        else:
            num = num258 * 3 + 2 
            if num > n:
                num = num369 * 3 + 3
                if num > n:
                    print(-1)
                    exit()
                ans[i] = num
                set3.remove(num)
                num369 += 1
                continue
            ans[i] = num # 2, 5, 8, 11, ...
            set2.remove(num)
            num258 += 1
     
    for i in range(n):
        if ans[i] == 0:
            if set1:
                ans[i] = set1.pop()
            elif set2:
                ans[i] = set2.pop()
            elif set3:
                ans[i] = set3.pop()
print(*ans)