Problem p02586 (python)

Fixed submission:
import sys
import math
import collections
import bisect
import itertools

# import numpy as np

sys.setrecursionlimit(10 ** 7)
INF = 10 ** 20
MOD = 10 ** 9 + 7
# MOD = 998244353

ni = lambda: int(sys.stdin.readline().rstrip())
ns = lambda: map(int, sys.stdin.readline().rstrip().split())
na = lambda: list(map(int, sys.stdin.readline().rstrip().split()))
na1 = lambda: list(map(lambda x: int(x) - 1, sys.stdin.readline().rstrip().split()))


# ===CODE===

def main():
    r, c, k = ns()

    mat = [[0 for _ in range(c + 1)] for __ in range(r + 1)]
    for _ in range(k):
        x, y, v = ns()
        x, y = x - 1, y - 1
        mat[x][y] = v

    dp = [[0] * 4 for _ in range(c + 1)]

    for i in range(r):
        ndp = [[0] * 4 for _ in range(c + 1)]
        for j in range(c):
            for n in range(2, -1, -1):
                dp[j][n + 1] = max(dp[j][n + 1], dp[j][n] + mat[i][j])

            ndp[j][0] = max(dp[j])
            for n in range(4):
                dp[j + 1][n] = max(dp[j + 1][n], dp[j][n])
        dp = ndp

    print(max(dp[c - 1]))


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
import sys
import math
import collections
import bisect
import itertools


import numpy as np
import numba
from numba import njit, b1, i4, i8, f8

sys.setrecursionlimit(10 ** 7)
INF = 10 ** 20
MOD = 10 ** 9 + 7
# MOD = 998244353

ni = lambda: int(sys.stdin.readline().rstrip())
ns = lambda: map(int, sys.stdin.readline().rstrip().split())
na = lambda: list(map(int, sys.stdin.readline().rstrip().split()))
na1 = lambda: list(map(lambda x: int(x) - 1, sys.stdin.readline().rstrip().split()))


# ===CODE===

@njit((i8, i8, i8, i8[:]), cache=True)
def main():
    r, c, k = ns()

    mat = [[0 for _ in range(c + 1)] for __ in range(r + 1)]
    for _ in range(k):
        x, y, v = ns()
        x, y = x - 1, y - 1
        mat[x][y] = v

    dp = [[0] * 4 for _ in range(c + 1)]

    for i in range(r):
        ndp = [[0] * 4 for _ in range(c + 1)]
        for j in range(c):
            for n in range(2, -1, -1):
                dp[j][n + 1] = max(dp[j][n + 1], dp[j][n] + mat[i][j])

            ndp[j][0] = max(dp[j])
            for n in range(4):
                dp[j + 1][n] = max(dp[j + 1][n], dp[j][n])
        dp = ndp

    print(max(dp[c - 1]))

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
import sys
import math
import collections
import bisect
import itertools
import numba
from numba import njit, b1, i4, i8, f8

# import numpy as np

sys.setrecursionlimit(10 ** 7)
INF = 10 ** 20
MOD = 10 ** 9 + 7
# MOD = 998244353

ni = lambda: int(sys.stdin.readline().rstrip())
ns = lambda: map(int, sys.stdin.readline().rstrip().split())
na = lambda: list(map(int, sys.stdin.readline().rstrip().split()))
na1 = lambda: list(map(lambda x: int(x) - 1, sys.stdin.readline().rstrip().split()))


# ===CODE===

@njit((i8, i8, i8, i8[:]), cache=True)
def main():
    r, c, k = ns()

    mat = [[0 for _ in range(c + 1)] for __ in range(r + 1)]
    for _ in range(k):
        x, y, v = ns()
        x, y = x - 1, y - 1
        mat[x][y] = v

    dp = [[0] * 4 for _ in range(c + 1)]

    for i in range(r):
        ndp = [[0] * 4 for _ in range(c + 1)]
        for j in range(c):
            for n in range(2, -1, -1):
                dp[j][n + 1] = max(dp[j][n + 1], dp[j][n] + mat[i][j])

            ndp[j][0] = max(dp[j])
            for n in range(4):
                dp[j + 1][n] = max(dp[j + 1][n], dp[j][n])
        dp = ndp

    print(max(dp[c - 1]))

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
import sys
import math
import collections
import bisect
import itertools
import numba
from numba import njit, b1, i4, i8, f8

# import numpy as np

sys.setrecursionlimit(10 ** 7)
INF = 10 ** 20
MOD = 10 ** 9 + 7
# MOD = 998244353

ni = lambda: int(sys.stdin.readline().rstrip())
ns = lambda: map(int, sys.stdin.readline().rstrip().split())
na = lambda: list(map(int, sys.stdin.readline().rstrip().split()))
na1 = lambda: list(map(lambda x: int(x) - 1, sys.stdin.readline().rstrip().split()))


# ===CODE===

@njit((i8, i8, i8, i8[:]), cache=True)
def main():
    r, c, k = ns()

    mat = [[0] * (c + 1)] * (r + 1)
    for _ in range(k):
        x, y, v = ns()
        x, y = x - 1, y - 1
        mat[x][y] = v

    dp = [[0] * 4] * (c + 1)

    for i in range(r):
        ndp = [[0] * 4] * (c + 1)
        for j in range(c):
            for n in range(2, -1, -1):
                dp[j][n + 1] = max(dp[j][n + 1], dp[j][n] + mat[i][j])

            ndp[j][0] = max(dp[j])
            for n in range(4):
                dp[j + 1][n] = max(dp[j + 1][n], dp[j][n])
        dp = ndp

    print(max(dp[c - 1]))


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
import sys
import math
import collections
import bisect
import itertools

# import numpy as np

sys.setrecursionlimit(10 ** 7)
INF = 10 ** 20
MOD = 10 ** 9 + 7
# MOD = 998244353

ni = lambda: int(sys.stdin.readline().rstrip())
ns = lambda: map(int, sys.stdin.readline().rstrip().split())
na = lambda: list(map(int, sys.stdin.readline().rstrip().split()))
na1 = lambda: list(map(lambda x: int(x) - 1, sys.stdin.readline().rstrip().split()))


# ===CODE===

# https://note.nkmk.me/python-union-find/

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())


def main():
    r, c, k = ns()

    mat = [[0 for _ in range(r)] for __ in range(c)]
    res = [[[0, 0, 0] for _ in range(r)] for __ in range(c)]
    visited = [[False for _ in range(r)] for __ in range(c)]
    visited[0][0] = True
    for _ in range(k):
        ri, ci, ki = ns()
        ri, ci = ri - 1, ci - 1
        mat[ci][ri] = ki

    que = collections.deque()
    tmp = [0, 0]
    que.append(tmp)

    while que:
        xi, yi = que.popleft()
        if yi + 1 < c:
            if not visited[yi + 1][xi]:
                que.append([yi + 1, xi])
                visited[yi + 1][xi] = True
            for i in range(2, -1, -1):
                if i == 0:
                    res[yi + 1][xi][i] = max(res[yi + 1][xi][i], mat[yi][xi])
                else:
                    if res[yi + 1][xi][i - 1] != 0:
                        res[yi + 1][xi][i] = max(res[yi + 1][xi][i], res[yi + 1][xi][i - 1] + mat[yi][xi])

        if xi + 1 < r:
            if not visited[yi][xi + 1]:
                que.append([yi, xi + 1])
                visited[yi][xi + 1] = True
            for i in range(2, -1, -1):
                if i == 0:
                    res[yi][xi + 1][i] = max(res[yi][xi + 1][i], mat[yi][xi])
                else:
                    if res[yi][xi + 1][i - 1]!= 0:
                        res[yi][xi + 1][i] = max(res[yi][xi + 1][i], res[yi][xi + 1][i - 1] + mat[yi][xi])

    print(max(res[c - 1][r - 1]))


if __name__ == '__main__':
    main()