Problem p03326 (python)

Fixed submission:
import sys
def input(): return sys.stdin.readline().strip()


def main():
    N, M = map(int, input().split())
    cakes = [(0, 0, 0)]
    for _ in range(N):
        x, y, z = map(int, input().split())
        cakes.append((x, y, z))

    """
    dp[n][m][t] = (1~n番目までのケーキからm個のケーキをタイプtで選んだ時の合計の最大値)
    と定める。ここでタイプとは
        t=0 -> val = x + y + z
        t=1 -> val = x + y - z
        t=2 -> val = x - y + z
        t=3 -> val = x - y - z
        t=4 -> val = -x + y + z
        t=5 -> val = -x + y - z
        t=6 -> val = -x - y + z
        t=7 -> val = -x - y - z
    どのタイプにせよ更新式は同じで、dp[n][m] = max(dp[n-1][m], dp[n-1][m-1] + (n番目のdp = [0]aaa))
    """

    dp = [[[0] * 8 for _ in range(M + 1)] for _ in range(N + 1)]
    for n in range(1, N + 1):
        x, y, z = cakes[n]
        for m in range(1, min(M + 1, n + 1)):
            if n == m:
                dp[n][m][0] = dp[n - 1][m - 1][0] + (x + y + z)
                dp[n][m][1] = dp[n - 1][m - 1][1] + (x + y - z)
                dp[n][m][2] = dp[n - 1][m - 1][2] + (x - y + z)
                dp[n][m][3] = dp[n - 1][m - 1][3] + (x - y - z)
                dp[n][m][4] = dp[n - 1][m - 1][4] + (-x + y + z)
                dp[n][m][5] = dp[n - 1][m - 1][5] + (-x + y - z)
                dp[n][m][6] = dp[n - 1][m - 1][6] + (-x - y + z)
                dp[n][m][7] = dp[n - 1][m - 1][7] + (-x - y - z)
            else:
                dp[n][m][0] = max(dp[n - 1][m][0], dp[n - 1][m - 1][0] + (x + y + z))
                dp[n][m][1] = max(dp[n - 1][m][1], dp[n - 1][m - 1][1] + (x + y - z))
                dp[n][m][2] = max(dp[n - 1][m][2], dp[n - 1][m - 1][2] + (x - y + z))
                dp[n][m][3] = max(dp[n - 1][m][3], dp[n - 1][m - 1][3] + (x - y - z))
                dp[n][m][4] = max(dp[n - 1][m][4], dp[n - 1][m - 1][4] + (-x + y + z))
                dp[n][m][5] = max(dp[n - 1][m][5], dp[n - 1][m - 1][5] + (-x + y - z))
                dp[n][m][6] = max(dp[n - 1][m][6], dp[n - 1][m - 1][6] + (-x - y + z))
                dp[n][m][7] = max(dp[n - 1][m][7], dp[n - 1][m - 1][7] + (-x - y - z))
    print(max(dp[N][M]))

if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys
def input(): return sys.stdin.readline().strip()


def quickSelect(A, k):
    """
    長さnの配列Aの中からk番目(1-indexed)に小さい値を線形時間で求めるアルゴリズム。
    動作原理は次を参照：http://www.flint.jp/blog/?entry=109
    処理の流れは次を参照：https://naoyat.hatenablog.jp/entry/median-of-medians

    使用例：quickSelect([2,7,1,8,2], 5, 3) => 2
    """
    n = len(A)
    while True:
        B = A
        m = n

        # 1.配列からmedian of mediansにより適切なpivotを選択（25%~75%に収まる）
        while m > 5:
            medians = []
            rem = 0
            for i in range(m // 5):
                sortB = sorted(B[i * 5: (i + 1) * 5])
                medians.append(sortB[2])
            if m % 5 != 0:
                sortB = sorted(B[(m // 5) * 5:])
                medians.append(sortB[(m - (m // 5) * 5) // 2])
                rem = 1
            B = medians
            m = m // 5 + rem
        B.sort()
        pivot = B[m // 2]

        # 2.pivotを軸に大小で左右に分ける
        left = []
        right = []
        left_num = 0
        right_num = 0
        for a in A:
            if a < pivot:
                left.append(a)
                left_num += 1
            elif a > pivot:
                right.append(a)
                right_num += 1
        if left_num < k <= n - right_num:
            return pivot
        elif left_num < k:
            A = right
            k -= (n - right_num)
            n = right_num
        else:
            A = left
            n = left_num


def main():
    N, M = map(int, input().split())
    cakes = [(0, 0, 0)]
    for _ in range(N):
        x, y, z = map(int, input().split())
        cakes.append((x, y, z))

    """
    quickSelectによりO(N)で解ける別解を実装。
    x+y+zやx-y-z等のそれぞれに対してcakesをソートして大きい値から取れば良い。
    """

    ans = 0
    for i in range(8):
        sorted_cakes = []
        for j in range(1, N + 1):
            S = 0
            for k in range(3): # iの立っているビットでx, y, zの符号を管理している！！
                if i & (1 << k):
                    S += cakes[j][k]
                else:
                    S -= cakes[j][k]
            sorted_cakes.append(S)

        # ここでsorted_cakesのM番目に大きな値（＝(N - M + 1)番目に小さな値）をquickSelectで求める
        # あとはその数以上なら評価に加算すれば良い
        val = 0
        bound = quickSelect(sorted_cakes, N - M + 1)
        cnt = 0
        for c in sorted_cakes:
            if c > bound:
                val += c
                cnt += 1
        val += bound * (M - cnt)
        ans = max(ans, val)
    print(ans)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys
def input(): return sys.stdin.readline().strip()


def main():
    N, M = map(int, input().split())
    cakes = [(0, 0, 0)]
    for _ in range(N):
        x, y, z = map(int, input().split())
        cakes.append((x, y, z))

    """
    dp[n][m][t] = (1~n番目までのケーキからm個のケーキをタイプtで選んだ時の合計の最大値)
    と定める。ここでタイプとは
        t=0 -> val = x + y + z
        t=1 -> val = x + y - z
        t=2 -> val = x - y + z
        t=3 -> val = x - y - z
        t=4 -> val = -x + y + z
        t=5 -> val = -x + y - z
        t=6 -> val = -x - y + z
        t=7 -> val = -x - y - z
    どのタイプにせよ更新式は同じで、dp[n][m] = max(dp[n-1][m], dp[n-1][m-1] + (n番目のdp = [0]aaa))
    """

    dp = [[[0] * 8 for _ in range(M + 1)] for _ in range(N + 1)]
    dp[0][1]
    for n in range(1, N + 1):
        x, y, z = cakes[n]
        for m in range(1, min(M + 1, n + 1)):
            if n == m:
                dp[n][m][0] = dp[n - 1][m - 1][0] + (x + y + z)
                dp[n][m][1] = dp[n - 1][m - 1][1] + (x + y - z)
                dp[n][m][2] = dp[n - 1][m - 1][2] + (x - y + z)
                dp[n][m][3] = dp[n - 1][m - 1][3] + (x - y - z)
                dp[n][m][4] = dp[n - 1][m - 1][4] + (-x + y + z)
                dp[n][m][5] = dp[n - 1][m - 1][5] + (-x + y - z)
                dp[n][m][6] = dp[n - 1][m - 1][6] + (-x - y + z)
                dp[n][m][7] = dp[n - 1][m - 1][7] + (-x - y - z)
            else:
                dp[n][m][0] = max(dp[n - 1][m][0], dp[n - 1][m - 1][0] + (x + y + z))
                dp[n][m][1] = max(dp[n - 1][m][1], dp[n - 1][m - 1][1] + (x + y - z))
                dp[n][m][2] = max(dp[n - 1][m][2], dp[n - 1][m - 1][2] + (x - y + z))
                dp[n][m][3] = max(dp[n - 1][m][3], dp[n - 1][m - 1][3] + (x - y - z))
                dp[n][m][4] = max(dp[n - 1][m][4], dp[n - 1][m - 1][4] + (-x + y + z))
                dp[n][m][5] = max(dp[n - 1][m][5], dp[n - 1][m - 1][5] + (-x + y - z))
                dp[n][m][6] = max(dp[n - 1][m][6], dp[n - 1][m - 1][6] + (-x - y + z))
                dp[n][m][7] = max(dp[n - 1][m][7], dp[n - 1][m - 1][7] + (-x - y - z))
    print(max(dp[N][M]))

if __name__ == "__main__":
    main()