Problem p02750 (Python)

Buggy submission (status: Time Limit Exceeded):
def main():
    # input
    N, T = map(int, input().split())
    a, b = [], []
    a0_b = []
    for i in range(N):
        a_in, b_in = map(int, input().split())
        if a_in == 0:
            a0_b.append(b_in)
        else:
            a.append(a_in)
            b.append(b_in)
    L = len(a)

    # sort
    ab_ratio = [a[i] / (b[i] + 1) for i in range(L)]
    key = sorted(range(L), key=lambda x: -ab_ratio[x])
    a = [a[i] for i in key]
    b = [b[i] for i in key]
    a0_b.sort()

    # --- case a >= 1 --- #
    # dp[for i shops][passed j shops] := min time
    # j:29 for T < 10^9
    J_MAX = min(L + 1, 29)
    dp = [[float("inf") for _ in range(J_MAX)] for _ in range(L + 1)]
    dp[0][0] = 0
    for i in range(L):
        for j in range(J_MAX):
            if j > 0:
                dt = 1 + a[i] * (dp[i][j - 1] + 1) + b[i]
                if dp[i][j - 1] + dt <= T:
                    dp[i + 1][j] = min(dp[i][j], dp[i][j - 1] + dt)
                else:
                    dp[i + 1][j] = dp[i][j]
            else:
                dp[i + 1][j] = dp[i][j]

    # --- case a = 0 --- #
    ans = 0
    for i in range(J_MAX):
        if dp[L][i] == float("inf"):
            break
        else:
            time_cnt = dp[L][i]
            shop_cnt = i
            for t_b in a0_b:
                if time_cnt + 1 + t_b <= T:
                    time_cnt += 1 + t_b
                    shop_cnt += 1
            if ans < shop_cnt:
                ans = shop_cnt

    print(ans)


if __name__ == "__main__":
    main()


Fixed submission:
import sys

input = sys.stdin.readline


def main():
    # input
    N, T = map(int, input().split())
    a, b = [], []
    a0_b = []
    for i in range(N):
        a_in, b_in = map(int, input().split())
        if a_in == 0:
            a0_b.append(b_in)
        else:
            a.append(a_in)
            b.append(b_in)
    L = len(a)

    # sort
    ab_ratio = [a[i] / (b[i] + 1) for i in range(L)]
    key = sorted(range(L), key=lambda x: -ab_ratio[x])
    a = [a[i] for i in key]
    b = [b[i] for i in key]
    a0_b.sort()

    # --- case a >= 1 --- #
    # dp[for i shops][passed j shops] := min time
    # j:29 for T < 10^9
    J_MAX = min(L + 1, 29)
    dp = [[float("inf") for _ in range(J_MAX)] for _ in range(L + 1)]
    dp[0][0] = 0
    for i in range(L):
        for j in range(J_MAX):
            if j > 0:
                dt = 1 + a[i] * (dp[i][j - 1] + 1) + b[i]
                if dp[i][j - 1] + dt <= T:
                    dp[i + 1][j] = min(dp[i][j], dp[i][j - 1] + dt)
                else:
                    dp[i + 1][j] = dp[i][j]
            else:
                dp[i + 1][j] = dp[i][j]

    # --- case a = 0 --- #
    ans = 0
    for i in range(J_MAX):
        if dp[L][i] == float("inf"):
            break
        else:
            time_cnt = dp[L][i]
            shop_cnt = i
            for t_b in a0_b:
                if time_cnt + 1 + t_b <= T:
                    time_cnt += 1 + t_b
                    shop_cnt += 1
            if ans < shop_cnt:
                ans = shop_cnt

    print(ans)


if __name__ == "__main__":
    main()