Нахождение K-го по величине элемента с помощью кучи

Я пытаюсь решить проблему с лит-кодом: kth-самый большой-элемент-в-массиве

Я знаю способ решить эту проблему — использовать кучу. Однако мне хотелось реализовать на практике свой собственный метод heapify, и вот мой код:

def findKthLargest(self, nums: List[int], k: int) -> int:
    
    def heapify(nums: List[int], i: int):
        print(nums, i)
        
        largest = i
        left = (2 * i) + 1
        right = (2 * i) + 2
        
        if left < len(nums) and nums[largest] < nums[left]:
            largest = left
        if right < len(nums) and nums[largest] < nums[right]:
            largest = right
        
        if largest != i:
            nums[i], nums[largest] = nums[largest], nums[i]
            print(nums)
            heapify(nums, largest)
    
    print(nums)
    for i in range(len(nums)-1, -1, -1):
        heapify(nums, i)
    print(nums)
    
    return nums[k-1]

Мой код в основном соответствует реализации, приведенной в редакционной статье:

    def max_heapify(heap_size, index):
        left, right = 2 * index + 1, 2 * index + 2
        largest = index
        if left < heap_size and lst[left] > lst[largest]:
            largest = left
        if right < heap_size and lst[right] > lst[largest]:
            largest = right
        if largest != index:
            lst[index], lst[largest] = lst[largest], lst[index]
            max_heapify(heap_size, largest)

    # heapify original lst
    for i in range(len(lst) // 2 - 1, -1, -1):
        max_heapify(len(lst), i)

И это сработало для 21/41 тестовых случаев и не работает. Ввод:

nums = [3,2,3,1,2,4,5,5,6]
k = 4

Мой код возвращает 3 вместо 4. Вот мой результат:

[3, 2, 3, 1, 2, 4, 5, 5, 6]
[3, 2, 3, 1, 2, 4, 5, 5, 6] 8
[3, 2, 3, 1, 2, 4, 5, 5, 6] 7
[3, 2, 3, 1, 2, 4, 5, 5, 6] 6
[3, 2, 3, 1, 2, 4, 5, 5, 6] 5
[3, 2, 3, 1, 2, 4, 5, 5, 6] 4
[3, 2, 3, 1, 2, 4, 5, 5, 6] 3
[3, 2, 3, 6, 2, 4, 5, 5, 1]
[3, 2, 3, 6, 2, 4, 5, 5, 1] 8
[3, 2, 3, 6, 2, 4, 5, 5, 1] 2
[3, 2, 5, 6, 2, 4, 3, 5, 1]
[3, 2, 5, 6, 2, 4, 3, 5, 1] 6
[3, 2, 5, 6, 2, 4, 3, 5, 1] 1
[3, 6, 5, 2, 2, 4, 3, 5, 1]
[3, 6, 5, 2, 2, 4, 3, 5, 1] 3
[3, 6, 5, 5, 2, 4, 3, 2, 1]
[3, 6, 5, 5, 2, 4, 3, 2, 1] 7
[3, 6, 5, 5, 2, 4, 3, 2, 1] 0
[6, 3, 5, 5, 2, 4, 3, 2, 1]
[6, 3, 5, 5, 2, 4, 3, 2, 1] 1
[6, 5, 5, 3, 2, 4, 3, 2, 1]
[6, 5, 5, 3, 2, 4, 3, 2, 1] 3
[6, 5, 5, 3, 2, 4, 3, 2, 1]

Я вижу, что 4 в индексе 5 никогда не сортируется после первых нескольких итераций. Почему это происходит? Что мне не хватает? Любая помощь будет оценена по достоинству.

🤔 А знаете ли вы, что...
В Python есть инструменты для создания графиков и визуализации данных, такие как библиотеки Matplotlib и Seaborn.


2
51
1

Ответ:

Решено

Есть несколько проблем с вашим (частичным) алгоритмом пирамидальной сортировки:

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

    • Начинание с последнего элемента не поможет правильно построить кучу, поскольку конечные узлы (которые включают в себя последние элементы) уже удовлетворяют свойству кучи и не нуждаются в создании кучи.
  2. После построения максимальной кучи самый большой элемент оказывается в корне. Затем мы:

    1. Поменяйте его местами с последним элементом в куче.

    2. Уменьшите размер кучи на единицу и восстановите ее, вызвав heapify на корневом узле.

    3. Повторите этот процесс k-1 раз.

После этого k-й по величине элемент находится в корне кучи:

def findKthLargest(self, nums: List[int], k: int) -> int:
    # added heap_size parameter
    def heapify(nums: List[int], i: int, heap_size: int):
        print(nums, i)

        largest = i
        left = (2 * i) + 1
        right = (2 * i) + 2

        if left < heap_size and nums[largest] < nums[left]:
            largest = left
        if right < heap_size and nums[largest] < nums[right]:
            largest = right

        if largest != i:
            nums[i], nums[largest] = nums[largest], nums[i]
            print(nums)
            heapify(nums, largest, heap_size)

    print(nums)
    # starting from the last non-leaf node
    for i in range(len(nums) // 2 - 1, -1, -1):
        heapify(nums, i, len(nums))
    print(nums)

    # performing k - 1 swaps
    for i in range(len(nums) - 1, len(nums) - k, -1):
        nums[0], nums[i] = nums[i], nums[0]
        heapify(nums, 0, i)  # heapify the reduced heap

    return nums[0]  # kth largest element is at the root