ROSALIND｜Partial Sort (PS)

發表於 更新於

給定長度為 n 的整數數列，求其中前 k 小的數字並且按遞增順序呈現。

Given: A positive integer $n \le 10^5$ and an array $A[1..n]$ of integers from $-10^5$ to $10^5$, a positive integer $k \le 1000$.

Return: The $k$ smallest elements of a sorted array $A$.

(https://rosalind.info/problems/ps/)

這題為 heap sort 的變體，它只要求提供部分排序的數列。解題的關鍵在於先從數列取出前 k 個數字建立 heap。這個 heap 的組成不一定滿足題目要求，所以我們要從數列剩餘的部分，挑出能讓 heap 最大值更小的數字來替換。

也就是說，逐一比對根節點與剩餘數字。如果根節點大於某個剩餘數字，表示當前 heap 的組成仍有調整空間，因此用該數字取代根節點，再重複 Floyd’s method 的操作來恢復 heap 的特性。重複這個步驟，直到檢查完所有剩餘數字，就能得到一個由前 k 小的數字組成的 heap。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def sift_down(l: list[int], i: int, n: int):
    candidates = [idx for idx in (i, 2 * i + 1, 2 * i + 2) if idx < n]
    largest = max(candidates, key=lambda x: l[x])
    if largest != i:
        l[largest], l[i] = l[i], l[largest]
        sift_down(l, largest ,n)
def heapify(l: list[int], k: int) -> list[int]:
    # 先取出前 k 個數字建立 heap，此時還不是最終型態
    for idx in range(k // 2 - 1, -1, -1):
        sift_down(l, idx, k)
    # 如果根節點大於剩餘的數字，表示 heap 的總和還有縮小的空間
    for idx in range(k, len(l)):
        if l[0] > l[idx]:
            l[0] = l[idx]
            sift_down(l, 0, k)
    return l

最後對這個大小為 k 的 heap 進行 heap sort，便能遞增排序數列前 k 小的數字。

1
2
3
4
5
6
def ps(l: list[int], k: int) -> list[int]:
    hp = heapify(l, k)
    for end in range(k - 1, 0, -1):
        hp[0], hp[end] = hp[end], hp[0]
        sift_down(hp, 0, end)
    return(hp[:k])