AOJ Lesson ALDS1 9_D “HeapSort”をpythonで解いてみた

AOJのLessonのHeapSortを解いた際の備忘録です。
「AOJにある解答例は難しすぎる、、ネットでもわかりやすい解説もない、、」
という状況だったので、素人である自分に合わせた可読性の高い(?)コードを残しておきます。

※ヒープソートを理解している前提で記事を書いています

問題

AOJ Lesson アルゴリズムとデータ構造 Heap Sort

N要素の数列Aが与えられます。最大ヒープを満たし、ヒープソートを行ったときに疑似コード25行目のmaxHeapifyにおけるスワップ回数の総数が最大となるような数列Aの順列を１つ出力してください。

疑似コードは問題ページで確認していただけると幸いです。

方針

max_heapifyメソッドがやっていることを改めて考察する

ヒープの再構築を行う際のメソッド。
ただし、この問題で注目しているmax_heapifyはソートの実行中のものであり、最大値を整列済みにした際に上に上がってしまった要素を適切な位置に戻すためのものなので、初期の構築のように、
for i in range(N//2,0 ,-1): で回す必要はなく、一回実行するだけでヒープが再構築される。

→ 最大値を整列済にする際のswap A[1], A[heap_size] = A[heap_size], A[1] で毎度、最小値が上に上がってしまう時、この問題で見ているswap回数が最も多くなる。

入力例1の場合は答えの一つに 23 9 15 2 5 3 12 1 があるが、この数列をheap sortする際に以下のように、sort中のmax_heapify 実行後にprint(A)を入れると、

from sys import stdin, maxsize
from typing import List

N = int(stdin.readline())
*A, = map(int, stdin.readline().split())
# 番兵の追加
A = [maxsize] + A

def max_heapify(A, index, heap_size):
    left_child_index = 2*index
    right_child_index = 2*index+1

    if left_child_index <= heap_size and A[left_child_index] > A[index]:
        largest = left_child_index
    else:
        largest = index
    if right_child_index <= heap_size and A[right_child_index] > A[largest]:
        largest = right_child_index

    # 子の方が値が小さい場合、swapする
    if largest != index:
        A[index], A[largest] = A[largest], A[index]
        # 再帰的に呼び出す(子nodeのindexがHより小さくなって終了する)
        max_heapify(A, largest, heap_size)

def heap_sort(A: List[int], N: int):
    heap_size = N

    # heapの構築
    for i in range(N//2, 0, -1):
        max_heapify(A, i, N)

    # sort
    while heap_size >= 2:
        # 1番目に格納されている最も大きい値を最後尾へ格納することを繰り返す
        A[1], A[heap_size] = A[heap_size], A[1]
        heap_size -= 1
        max_heapify(A, 1, heap_size)
                # 整列の状況を出力する
        print(A)

heap_sort(A, N)

出力

[9223372036854775807, 15, 9, 12, 2, 5, 3, 1, 23]
[9223372036854775807, 12, 9, 3, 2, 5, 1, 15, 23]
[9223372036854775807, 9, 5, 3, 2, 1, 12, 15, 23]
[9223372036854775807, 5, 2, 3, 1, 9, 12, 15, 23]
[9223372036854775807, 3, 2, 1, 5, 9, 12, 15, 23]
[9223372036854775807, 2, 1, 3, 5, 9, 12, 15, 23]
[9223372036854775807, 1, 2, 3, 5, 9, 12, 15, 23]

と、整列済のひとつ前に必ず最小値である1が来るようになっている！！
つまり、この問題は、ソート中のmax_heapify後に未整列部の最小値が未整列部の最後尾に来るような数列を作成することに帰着します。
この事実に気づくまでにかなりの時間をかけてしまいました😅

解答の作成

ヒープソートでの整列はheap_sizeがNから2にデクリメントしながら行われる
⇒ 今回の解答では 2からNにインクリメントしながらヒープソートの逆順で目的の数列を作成する

1. 1番目 に最小値、i番目に対象としている範囲の最大値が格納されているので、 ⇒まず A[1]とA[i]をスワップ
2. ヒープの下に行ってしまった最小値をもう一度上へ引き戻す
3. 1,2をN番目まで繰り返す

最後に再度、1番目の最小値、対象としている範囲の最大値であるN番目の値をスワップして終了する

実装例

# ...上のコードの続き...

heap_sort(A, N)

for i in range(2, N):
    A[1], A[i] = A[i], A[1]
    # while文を用いて最も小さい値をrootまで戻す
    while i != 1:
        A[i], A[i//2] = A[i//2], A[i]
        i //= 2
A[1], A[N] = A[N], A[1]

print(*A[1:])

最後に

ヒープソートを理解したばかりの状態でこの問題を考えるのは、正直苦労しました。
間違えなどあればコメントいただけると幸いです。