1329. Sort the Matrix Diagonally

Python

Build Map, Sort, Put back

N是row*col的matrix總數

• Time: O(NlogN) # sorting
• Space: O(N)

from collections import defaultdict, deque


class Solution:
    def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        mapper = defaultdict(list)

        # Build the map
        for row in range(m):
            for col in range(n):
                mapper[row-col].append(mat[row][col])

        # Sort them
        for id in mapper.keys():
            mapper[id].sort()
            mapper[id] = deque(mapper[id])

        # Put back
        for row in range(m):
            for col in range(n):
                mat[row][col] = mapper[row-col].popleft()
        return mat

Build with Heap

• 少一次loop，build的時候就排好
• 會快一點點

import heapq
from collections import defaultdict


class Solution:
    def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        mapper = defaultdict(list)

        for row in range(m):
            for col in range(n):
                heapq.heappush(mapper[row-col], mat[row][col])

        for row in range(m):
            for col in range(n):
                mat[row][col] = heapq.heappop(mapper[row-col])

        return mat