Задача: 1329. Sort the Matrix Diagonally
Сложность: medium

Диагональ матрицы — это диагональная линия ячеек, начинающаяся с какой-либо ячейки в самой верхней строке или в самом левом столбце и идущая в направлении вниз-вправо до конца матрицы. Например, диагональ матрицы, начинающаяся с mat[2][0], где mat — это матрица размером 6 x 3, включает ячейки mat[2][0], mat[3][1] и mat[4][2].

Дана матрица mat размером m x n, состоящая из целых чисел. Отсортируйте каждую диагональ матрицы по возрастанию и верните полученную матрицу.

Пример:

Input: mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]
Output: [[1,1,1,1],[1,2,2,2],[1,2,3,3]]

👨‍💻 Алгоритм:

1⃣Сохраните размеры матрицы m и n. Создайте хеш-карту из минимальных куч для хранения элементов диагоналей.

2⃣Вставьте значения в хеш-карту, используя разность между индексами строки и столбца как ключ, чтобы собирать элементы на одной и той же диагонали.

3⃣Извлеките значения из хеш-карты и обновите матрицу, заполняя ее отсортированными значениями диагоналей. Верните отсортированную матрицу.

😎 Решение

class Solution {
    func diagonalSort(_ mat: [[Int]]) -> [[Int]] {
        var mat = mat
        let m = mat.count
        let n = mat[0].count
        
        var diagonals = [Int: PriorityQueue<Int>]()
        
        for row in 0..<m {
            for col in 0..<n {
                let key = row - col
                if diagonals[key] == nil {
                    diagonals[key] = PriorityQueue<Int>(order: <)
                }
                diagonals[key]?.push(mat[row][col])
            }
        }
        
        for row in 0..<m {
            for col in 0..<n {
                let key = row - col
                mat[row][col] = diagonals[key]?.pop() ?? 0
            }
        }
        
        return mat
    }
}

struct PriorityQueue<T: Comparable> {
    private var heap: [T]
    private let order: (T, T) -> Bool

    init(order: @escaping (T, T) -> Bool) {
        self.heap = []
        self.order = order
    }

    var isEmpty: Bool { return heap.isEmpty }
    var count: Int { return heap.count }

    mutating func push(_ element: T) {
        heap.append(element)
        siftUp(heap.count - 1)
    }

    mutating func pop() -> T? {
        guard !heap.isEmpty else { return nil }
        if heap.count == 1 {
            return heap.removeFirst()
        } else {
            let value = heap[0]
            heap[0] = heap.removeLast()
            siftDown(0)
            return value
        }
    }

    private mutating func siftUp(_ index: Int) {
        var childIndex = index
        let child = heap[childIndex]
        var parentIndex = (childIndex - 1) / 2
        
        while childIndex > 0 && order(child, heap[parentIndex]) {
            heap[childIndex] = heap[parentIndex]
            childIndex = parentIndex
            parentIndex = (childIndex - 1) / 2
        }
        heap[childIndex] = child
    }

    private mutating func siftDown(_ index: Int) {
        var parentIndex = index
        let count = heap.count
        let element = heap[parentIndex]
        var childIndex = (parentIndex * 2) + 1

        while childIndex < count {
            let rightChildIndex = childIndex + 1
            if rightChildIndex < count && order(heap[rightChildIndex], heap[childIndex]) {
                childIndex = rightChildIndex
            }
            if order(heap[childIndex], element) {
                heap[parentIndex] = heap[childIndex]
                parentIndex = childIndex
                childIndex = (parentIndex * 2) + 1
            } else {
                break
            }
        }
        heap[parentIndex] = element
    }
}

Ставь 👍 и забирай 📚 Базу знаний