Problem statement

https://leetcode.com/problems/get-biggest-three-rhombus-sums-in-a-grid/

Solution

If we want to achieve complexity better than O(m^2n^2), we need to cache some values.

d1 is used to hash one type of diagonals, which goes in direction (1, 1). Imagine, than we have original grid:

01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 17 18 19 20

Then table dp(i, j, -1) will have values:

01 02 03 04 05 06 08 10 12 14 11 18 21 24 27 16 28 36 40 44

dp(i, j, 1) is used to hash another type of diagonals. It will look like this:

01 08 21 40 44 06 18 36 39 42 11 28 30 32 34 16 17 18 19 20

We will use lru_cache options here to get fast updates. Then we need to traverse all possible rhombus: we will have O(mn*min(m,n)) of them and carefully evaluate sum of elements on border: which can be separated into 4 sides, which I denoted p1, p2, p3, p4. Actually I do not like which indexes I have here, I think it can be improved (UPD, it is now)

Complexity

Time complexity is O(m*n*min(m,n)), space complexity is O(mn).

Code

class Solution:
    def getBiggestThree(self, grid):
        m, n, heap = len(grid), len(grid[0]), []
        
        def update(heap, num):
            if num not in heap:
                heappush(heap, num)
                if len(heap) > 3: heappop(heap)
            return heap
        
        for num in chain(*grid): update(heap, num)
          
        @lru_cache(None)
        def dp(i, j, dr):
            if not 0 <= i < n or not 0 <= j < m: return 0
            return dp(i-1, j+dr, dr) + grid[j][i]
        
        for q in range(1, (1 + min(m, n))//2):
            for i in range(q, n - q):
                for j in range(q, m - q):
                    p1 = dp(i + q, j, -1) - dp(i, j - q, -1)
                    p2 = dp(i - 1, j + q - 1, -1) - dp(i - q - 1, j - 1, -1)
                    p3 = dp(i, j - q, 1) - dp(i - q, j, 1)
                    p4 = dp(i + q - 1, j + 1, 1) - dp(i - 1, j + q + 1, 1)
                    update(heap, p1 + p2 + p3 + p4)

        return sorted(heap)[::-1]