Задача: 378. Kth Smallest Element in a Sorted Matrix
Сложность: medium

Дана матрица размером n x n, где каждая строка и каждый столбец отсортированы в порядке возрастания. Верните k-й наименьший элемент в матрице.

Заметьте, что это k-й наименьший элемент в отсортированном порядке, а не k-й уникальный элемент.

Вы должны найти решение с использованием памяти лучше, чем O(n²).

Пример:

Input: matrix = [[-5]], k = 1
Output: -5

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

1⃣Инициализировать мин-кучу H. В нашем решении мы будем рассматривать каждую строку как отдельный список. Поскольку столбцы также отсортированы, мы можем рассматривать каждый столбец как отдельный список.

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

3⃣Мин-куча будет содержать тройки информации (значение, строка, столбец). Куча будет упорядочена по значениям, и мы будем использовать номера строк и столбцов для добавления следующего элемента, если текущий элемент будет удален из кучи.

😎 Решение:

class MyHeapNode {
    var row: Int
    var column: Int
    var value: Int

    init(_ value: Int, _ row: Int, _ column: Int) {
        self.value = value
        self.row = row
        self.column = column
    }
}

class Solution {
    func kthSmallest(_ matrix: [[Int]], _ k: Int) -> Int {
        let N = matrix.count
        var minHeap = PriorityQueue<MyHeapNode> { $0.value < $1.value }

        for r in 0..<min(N, k) {
            minHeap.enqueue(MyHeapNode(matrix[r][0], r, 0))
        }

        var element: MyHeapNode = minHeap.peek()!
        var k = k
        while k > 0 {
            element = minHeap.dequeue()!
            let r = element.row, c = element.column

            if c < N - 1 {
                minHeap.enqueue(MyHeapNode(matrix[r][c + 1], r, c + 1))
            }
            k -= 1
        }

        return element.value
    }
}

struct PriorityQueue<Element: Equatable> {
    private var elements: [Element]
    private let priorityFunction: (Element, Element) -> Bool

    init(priorityFunction: @escaping (Element, Element) -> Bool) {
        self.elements = []
        self.priorityFunction = priorityFunction
    }

    var isEmpty: Bool {
        return elements.isEmpty
    }

    func peek() -> Element? {
        return elements.first
    }

    mutating func enqueue(_ element: Element) {
        elements.append(element)
        siftUp(elementAtIndex: elements.count - 1)
    }

    mutating func dequeue() -> Element? {
        guard !isEmpty else { return nil }
        if elements.count == 1 {
            return elements.removeLast()
        } else {
            let value = elements[0]
            elements[0] = elements.removeLast()
            siftDown(elementAtIndex: 0)
            return value
        }
    }

    private mutating func siftUp(elementAtIndex index: Int) {
        let parent = parentIndex(of: index)
        guard !isRoot(index), isHigherPriority(at: index, than: parent) else { return }
        swapElement(at: index, with: parent)
        siftUp(elementAtIndex: parent)
    }

    private mutating func siftDown(elementAtIndex index: Int) {
        let childIndex = highestPriorityIndex(for: index)
        if index == childIndex { return }
        swapElement(at: index, with: childIndex)
        siftDown(elementAtIndex: childIndex)
    }

    private func isRoot(_ index: Int) -> Bool {
        return index == 0
    }

    private func leftChildIndex(of index: Int) -> Int {
        return 2 * index + 1
    }

    private func rightChildIndex(of index: Int) -> Int {
        return 2 * index + 2
    }

    private func parentIndex(of index: Int) -> Int {
        return (index - 1) / 2
    }

    private func isHigherPriority(at firstIndex: Int, than secondIndex: Int) -> Bool {
        return priorityFunction(elements[firstIndex], elements[secondIndex])
    }

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