leetcode1289

class Solution:
    def find2Smallest(self,arr):
        n = len(arr)
        minsum1,minsum2 = sys.maxsize,sys.maxsize
        minindex1,minindex2 = -1,-1
        for j in range(n):
            if arr[j] <= minsum1 and arr[j] <= minsum2:
                minsum2 = minsum1
                minindex2 = minindex1
                minsum1 = arr[j]
                minindex1 = j
            elif arr[j] >= minsum1 and arr[j] <= minsum2:
                minsum2 = arr[j]
                minindex2 = j

        return [[minsum1,minindex1],[minsum2,minindex2]]

    def minFallingPathSum(self, arr: 'List[List[int]]') -> int:
        n = len(arr)
        if n == 1:
            return arr[0][0]

        for i in range(1,n):
            dp = self.find2Smallest(arr[i-1])
            for j in range(n):
                if j == dp[0][1]:
                    arr[i][j] += dp[1][0]
                else:
                    arr[i][j] += dp[0][0]
        return min(arr[-1])

本题是leetcode 931的升级版，如果直接计算每一行和其上面每一行中不在同一列的和，会TLE。因此只保留上一行最小的2个数字。

在计算当前行时，只需要判断当前单元格的列是否与前一行的最小值在同一列，如果在同一列，则使用倒数第二小的值。否则就使用上一行的最小值。

下面是TLE的方案，可以作为对比：

import sys
class Solution:
    def minFallingPathSum(self, arr: 'List[List[int]]') -> int:
        n = len(arr)
        if n == 1:
            return arr[0][0]
        dp = [[sys.maxsize for _ in range(n)]for _ in range(n)]
        dp[0] = arr[0]
        minsum = sys.maxsize
        for i in range(1,n):
            for j in range(n):
                for k in range(n):
                    if j != k:
                        dp[i][j] = min(dp[i][j],dp[i-1][k]+arr[i][j])
                        if i == n - 1:
                            minsum = min(minsum,dp[i][j])
        return minsum