1 class Solution: 2 def find2Smallest(self,arr): 3 n = len(arr) 4 minsum1,minsum2 = sys.maxsize,sys.maxsize 5 minindex1,minindex2 = -1,-1 6 for j in range(n): 7 if arr[j] <= minsum1 and arr[j] <= minsum2: 8 minsum2 = minsum1 9 minindex2 = minindex1 10 minsum1 = arr[j] 11 minindex1 = j 12 elif arr[j] >= minsum1 and arr[j] <= minsum2: 13 minsum2 = arr[j] 14 minindex2 = j 15 16 return [[minsum1,minindex1],[minsum2,minindex2]] 17 18 def minFallingPathSum(self, arr: 'List[List[int]]') -> int: 19 n = len(arr) 20 if n == 1: 21 return arr[0][0] 22 23 for i in range(1,n): 24 dp = self.find2Smallest(arr[i-1]) 25 for j in range(n): 26 if j == dp[0][1]: 27 arr[i][j] += dp[1][0] 28 else: 29 arr[i][j] += dp[0][0] 30 return min(arr[-1])
本题是leetcode 931的升级版，如果直接计算每一行和其上面每一行中不在同一列的和，会TLE。因此只保留上一行最小的2个数字。
在计算当前行时，只需要判断当前单元格的列是否与前一行的最小值在同一列，如果在同一列，则使用倒数第二小的值。否则就使用上一行的最小值。
下面是TLE的方案，可以作为对比：
1 import sys 2 class Solution: 3 def minFallingPathSum(self, arr: 'List[List[int]]') -> int: 4 n = len(arr) 5 if n == 1: 6 return arr[0][0] 7 dp = [[sys.maxsize for _ in range(n)]for _ in range(n)] 8 dp[0] = arr[0] 9 minsum = sys.maxsize 10 for i in range(1,n): 11 for j in range(n): 12 for k in range(n): 13 if j != k: 14 dp[i][j] = min(dp[i][j],dp[i-1][k]+arr[i][j]) 15 if i == n - 1: 16 minsum = min(minsum,dp[i][j]) 17 return minsum
