【Python】マージソートのコード

はじめに

参考文献にある書籍には、マージソートの擬似コードが書かれています。これを Python で書き換えたコードを紹介する記事です。

マージソートのコード

マージソートの擬似コードを Python で書き換えると次のようになります。

マージソートのコード

import sys

sys.setrecursionlimit(10 ** 9) # 再帰回数の上限を設定

def merge_sort(array, left_index, right_index):
    if left_index < right_index:
        mid = (left_index + right_index) // 2
        merge_sort(array, left_index, mid)
        merge_sort(array, mid + 1, right_index)
        merge(array, left_index, mid, right_index)
    return array

def merge(array, left_index, mid, right_index):
    array1_size = mid - left_index + 1
    array2_size = right_index - mid
    array1 = []
    array2 = []
    for i in range(array1_size):
        array1.append(array[left_index + i])
    for j in range(array2_size):
        array2.append(array[mid + j + 1])
    array1.append(float('inf'))
    array2.append(float('inf'))
    i = 0
    j = 0
    for k in range(left_index, right_index + 1):
        if array1[i] <= array2[j]:
            array[k] = array1[i]
            i += 1
        else:
            array[k] = array2[j]
            j += 1
    return array

if __name__ == '__main__':
    array = [5, 3, 1, 6, 4, 2]
    result = merge_sort(array, 0, len(array) - 1)
    print(result) # [1, 2, 3, 4, 5, 6]

マージソートのコード（改良版）

関数 merge_sort() を呼び出すときに、配列全体をソートする場合は第二引数および第三引数を省略できるようにしたいです。また、場合によっては降順にソートできると嬉しいです。
これらの機能を実現するために改良した関数 merge_sort_main() の第二引数に reverse を追加することで True の場合は降順にソートし、 False の場合は昇順にソートするよう指定できるようにしました。デフォルトは False であるので昇順にソートされるのがデフォルトとなります。

マージソートのコード（改良版）

import sys

sys.setrecursionlimit(10 ** 9) # 再帰回数の上限を設定

def merge_sort_main(a, reverse=False):
    array = a.copy()
    result = merge_sort(array, 0, len(array) - 1, reverse)
    return result

def merge_sort(array, left_index, right_index, reverse):
    if left_index < right_index:
        mid = (left_index + right_index) // 2
        merge_sort(array, left_index, mid, reverse)
        merge_sort(array, mid + 1, right_index, reverse)
        merge(array, left_index, mid, right_index, reverse)
    return array

def merge(array, left_index, mid, right_index, reverse):
    array1_size = mid - left_index + 1
    array2_size = right_index - mid
    array1 = []
    array2 = []
    for i in range(array1_size):
        array1.append(array[left_index + i])
    for j in range(array2_size):
        array2.append(array[mid + j + 1])
    if reverse:
        array1.append(-float('inf'))
        array2.append(-float('inf'))
    else:
        array1.append(float('inf'))
        array2.append(float('inf'))
    i = 0
    j = 0
    for k in range(left_index, right_index + 1):
        if reverse:
            if array1[i] >= array2[j]:
                array[k] = array1[i]
                i += 1
            else:
                array[k] = array2[j]
                j += 1
        else:
            if array1[i] <= array2[j]:
                array[k] = array1[i]
                i += 1
            else:
                array[k] = array2[j]
                j += 1
    return array

if __name__ == '__main__':
    array = [5, 3, 1, 6, 4, 2]
    result1 = merge_sort_main(array)
    result2 = merge_sort_main(array, reverse=True)
    print(result1) # [1, 2, 3, 4, 5, 6]
    print(result2) # [6, 5, 4, 3, 2, 1]

参考文献

• コルメン, T., ライザーソン, C., リベスト, R., シュタイン, C. (著), 浅野 哲夫, 岩野 和生, 梅雄 博司, 山下 将史, 和田 幸一 (訳). 『世界標準MIT教科書 アルゴリズムイントロダクション 第3版 総合版』. 近代科学社 (2013).