Pythonでマージソートをなるべく少ない行数で実装してみた

友人がときどきアルゴリズムの実装の勉強していて、マージソートの実装に取り組んでいると聞いて俺もやるか〜ってなるのは多分きっと自然な流れ。
せっかくだからなんか勝負しようぜってことでなるべく少ない行数で実装したほうが勝ちみたいな謎勝負が行われました。
言語は友人がPythonでやろうと思ってたとのことだったのでPythonでやることになりました。負けないぞぉ

ルール

• データ配列は1〜10までの数字がランダムに並んだものを使う
• データ配列の生成、出力部分のコードは共通として、いじらないようにする
• とにかく[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]が出力されれば良い

共通部分のコードは以下の通り

from data import generate_data

data = generate_data() # 1〜10の数値がバラバラに並んでる
print(data)

# この下に実装

別ファイルにdata.pyがあり、そこのgenerate_data関数でデータを生成します。
データは[5, 1, 3, 9, 2, 8, 7, 4, 10, 6]みたいなのが生成されます。
data.pyの中身は適当になんとなく作っただけだし、ここでは省略します。

実装結果

上記のルールに則ってゴリゴリに書いたらこんな感じになりました。

from data import generate_data

data = generate_data() # 1〜10の数値がバラバラに並んでる
print(data)

# この下に実装
divide = lambda data: data if len(data) == 1 else [min(data), max(data)] if len(data) <= 2 else [divide([data[i] for i in range(len(data) // 2)]), divide([data[i] for i in range(len(data) // 2, len(data))])]
merge = lambda left, right: [right.pop(0) if len(left) == 0 else left.pop(0) if len(right) == 0 else left.pop(0) if left[0] < right[0] else right.pop(0) for i in range(len(left) + len(right))] if type(left[0]) is int and type(right[0]) is int else merge(left, merge(right[0], right[1])) if type(left[0]) is int else (merge(merge(left[0], left[1]), right) if type(right[0]) is int else merge(merge(left[0], left[1]), merge(right[0], right[1])))
print(merge(divide(data)[0], divide(data)[1]))

出力

[5, 1, 3, 9, 2, 8, 7, 4, 10, 6]
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

「長い。三行で」と言われても問題ない出来になりました。とりあえず解説していきます。

おおまかな流れ

マージソートは「分割統治法」というやつに基づいているらしく、一旦データ配列を最小単位まで分割し、分割した各データから適切にデータを選んで結合しながら元に戻して並び替えするらしいです。考えた人すげー
アルゴリズムの詳細については数多の先人たちが色んな解説をしているのでここでは省きます。

というわけで、上記の3行の実装結果は
1行目: 分割
2行目: 結合
3行目: 出力
といった感じに役割分担されてます。

まずは「分割」から見ていきます。

「分割」の実装

目的

該当のコードはこれですね。

「分割」のコード

divide = lambda data: data if len(data) == 1 else [min(data), max(data)] if len(data) <= 2 else [divide([data[i] for i in range(len(data) // 2)]), divide([data[i] for i in range(len(data) // 2, len(data))])]

データ配列を[7, 6, 3, 2, 4, 5, 1, 8]とすると「分割」の目的はこの配列から

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

のような階層構造を持つ多重配列を生成するということになります。
最小単位を要素数が2以下になるように半分ずつに分割していき、最小単位の要素数が2以下だったら小さい順にしておく
みたいなことをやってます。小さい順にするやつは本来「結合」の役目ですが、ここでやっておくと都合が良いのでこの時点でやっちゃいます。

書き換え

「分割」の関数は本来こんな感じになってます。

本来の形

def divide(data):
    if (len(data) <= 2):
        return data if len(data) == 1 else [min(data), max[data]]
    else:
        half_count = len(data) // 2
        left = [data[i] for i in range(half_count)]
        right = [data[i] for i in range(half_count, len(data))]
        return [divide(left), divide(right)]

要素数が2以下のときは小さい順にした配列を返し、それ以外のときはleftとrightに半分ずつ分けて再帰することで上記の階層構造を実現することができます。

else以降、half_countの変数使っているところは直接len(data) // 2に置き換えれば変数を使わずに済みます。あ、leftとrightも一回しか使ってないじゃん！そのまま返り値内にぶちこんじゃえ〜ってやると

返り値にぶちこんじゃえ〜

def divide(data):
    if (len(data) <= 2):
        return data if len(data) == 1 else [min(data), max[data]]
    else:
        return [divide([data[i] for i in range(len(data) // 2)]), divide([data[i] for i in range(len(data) // 2, len(data))])]

になります。
あ、ifとelseのどっちもreturn1文だけじゃん。それ三項演算子で1行じゃんってことで

三項演算子で1行じゃん

def divide(data):
    return data if len(data) == 1 else [min(data), max(data)] if len(data) <= 2 else [divide([data[i] for i in range(len(data) // 2)]), divide([data[i] for i in range(len(data) // 2, len(data))])]

といって感じに落ち着きます。関数内にreturnが1文しかないならラムダ式使って

ラムダ式にしちゃえ

divide = lambda data: data if len(data) == 1 else [min(data), max(data)] if len(data) <= 2 else [divide([data[i] for i in range(len(data) // 2)]), divide([data[i] for i in range(len(data) // 2, len(data))])]

見事1行に収まりました。ラムダ式でも再帰が使えるのはいいですね。

「結合」の実装

目的

該当のコードはこれですね。

「結合」のコード

merge = lambda left, right: [right.pop(0) if len(left) == 0 else left.pop(0) if len(right) == 0 else left.pop(0) if left[0] < right[0] else right.pop(0) for i in range(len(left) + len(right))] if type(left[0]) is int and type(right[0]) is int else merge(left, merge(right[0], right[1])) if type(left[0]) is int else (merge(merge(left[0], left[1]), right) if type(right[0]) is int else merge(merge(left[0], left[1]), merge(right[0], right[1])))

ぐへぇ、なげぇよ...
「結合」の目的は

分割後のデータ

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

の分割後のデータを下の階層から順に結合していき、最終的に全て結合して階層1つの元の形に戻すことです。
流れとしては

流れ1

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

流れ2

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

といった具合です。この流れを実現するために、merge関数では

• 引数の左右の要素が両方ともその子要素が数値のときに結合する
• 結合は左右の最初の要素の小さいほうをpopし、新しい配列の子要素として加えて行う
• それ以外のときは左右の要素の子要素で再帰し、左右の要素の子要素が数値になるようにする

といったことをやります。言葉で説明しても理解しにくいのでコードを見ましょう。

書き換え

「結合」の関数は本来こんな感じになってます。

本来の形

def merge(left, right):
    if type(left[0]) is int and type(right[0]) is int: # 両方とも子要素が数値のとき
        result = [] # 新しい配列を用意する
        for i in range(len(left) + len(right)):
            # left, rightの要素が無くなったらまだあるほうをpopしていく
            if len(left) == 0:
                result.append(right.pop(0))
            elif len(right) == 0:
                result.append(left.pop(0))
            # どっちもまだ残ってたら小さいほうをpopする
            elif left[0] < right[0]:
                result.append(left.pop(0))
            else:
                result.append(right.pop(0))
        return result
    else:
        """ 
        階層によっては同じレベルでも子要素が数値だったり配列だったりするので
        子要素が数値のやつはそのままにして再帰する
        """
        if type(left[0]) is int:
            return merge(left, merge(right[0], right[1]))
        elif type(right[0]) is int:
            return merge(merge(left[0], left[1]), right)
        else:
            return merge(merge(left[0], left[1]), merge(right[0], right[1]))

実装の詳細はコメントにて。
わざとif-elif-elseで完結するようにしました。ということは「分割」のときと同様に、三項演算子で返り値を1行にすることができて、さらにラムダ式で関数を1行で記述することができますね！

resultをリスト内包表記で表そうとしたとき、popでちゃんと要素減ってくれんのかなと心配していましたが、実際にやってみたらちゃんと減っててくれたのでリスト内包表記と三項演算子を使ってforの部分を1行で表すことができました。

「出力」の実装

「出力」のコード

print(merge(divide(data)[0], divide(data)[1]))

ここ本来は

data = divide(data)
print(merge(data[0], data[1]))

にしてdivideの実行は1回に留めるべきですが、ご覧のように変数の確保で1行使ってしまって行数の無駄になるので無理やり二度実行してます。まあ要素数10個程度じゃ誤差だよ誤差

おわり

間違いなく俺の中のクソコード・オブ・ザ・イヤー大賞の受賞コードになると思われます。ありがとうございました。