leetcode1289

 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