チラシの裏

プログラミングとか色々

OCaml標準ライブラリのList.sortを読む

 この記事はML Advent Calendar 2017の10日目(!?)のために書かれました.

 皆さんOCaml標準ライブラリのList.sortは使っていますか? 頻繁にリストに対してのソートを行うようならデータ構造かアルゴリズムを考え直した方が良いでしょうが,競プロの様にちょっとしたプログラムを書く時には僕もよくお世話になります. マニュアルによると,そのList.sortマージソートで実装されているようですが,僕のような者がちょっとやそっと工夫した程度では敵わないような最適化が施されています.今回はそのソースコードを読んでみましょう.

 まず,OCaml 4.06.0のlist.mlから引っぱってきたList.sortソースコードを見ていきましょう.

(** sorting *)

let rec merge cmp l1 l2 =
  match l1, l2 with
  | [], l2 -> l2
  | l1, [] -> l1
  | h1 :: t1, h2 :: t2 ->
      if cmp h1 h2 <= 0
      then h1 :: merge cmp t1 l2
      else h2 :: merge cmp l1 t2


let rec chop k l =
  if k = 0 then l else begin
    match l with
    | _::t -> chop (k-1) t
    | _ -> assert false
  end


let stable_sort cmp l =
  let rec rev_merge l1 l2 accu =
    match l1, l2 with
    | [], l2 -> rev_append l2 accu
    | l1, [] -> rev_append l1 accu
    | h1::t1, h2::t2 ->
        if cmp h1 h2 <= 0
        then rev_merge t1 l2 (h1::accu)
        else rev_merge l1 t2 (h2::accu)
  in
  let rec rev_merge_rev l1 l2 accu =
    match l1, l2 with
    | [], l2 -> rev_append l2 accu
    | l1, [] -> rev_append l1 accu
    | h1::t1, h2::t2 ->
        if cmp h1 h2 > 0
        then rev_merge_rev t1 l2 (h1::accu)
        else rev_merge_rev l1 t2 (h2::accu)
  in
  let rec sort n l =
    match n, l with
    | 2, x1 :: x2 :: _ ->
       if cmp x1 x2 <= 0 then [x1; x2] else [x2; x1]
    | 3, x1 :: x2 :: x3 :: _ ->
       if cmp x1 x2 <= 0 then begin
         if cmp x2 x3 <= 0 then [x1; x2; x3]
         else if cmp x1 x3 <= 0 then [x1; x3; x2]
         else [x3; x1; x2]
       end else begin
         if cmp x1 x3 <= 0 then [x2; x1; x3]
         else if cmp x2 x3 <= 0 then [x2; x3; x1]
         else [x3; x2; x1]
       end
    | n, l ->
       let n1 = n asr 1 in
       let n2 = n - n1 in
       let l2 = chop n1 l in
       let s1 = rev_sort n1 l in
       let s2 = rev_sort n2 l2 in
       rev_merge_rev s1 s2 []
  and rev_sort n l =
    match n, l with
    | 2, x1 :: x2 :: _ ->
       if cmp x1 x2 > 0 then [x1; x2] else [x2; x1]
    | 3, x1 :: x2 :: x3 :: _ ->
       if cmp x1 x2 > 0 then begin
         if cmp x2 x3 > 0 then [x1; x2; x3]
         else if cmp x1 x3 > 0 then [x1; x3; x2]
         else [x3; x1; x2]
       end else begin
         if cmp x1 x3 > 0 then [x2; x1; x3]
         else if cmp x2 x3 > 0 then [x2; x3; x1]
         else [x3; x2; x1]
       end
    | n, l ->
       let n1 = n asr 1 in
       let n2 = n - n1 in
       let l2 = chop n1 l in
       let s1 = sort n1 l in
       let s2 = sort n2 l2 in
       rev_merge s1 s2 []
  in
  let len = length l in
  if len < 2 then l else sort len l


let sort = stable_sort
let fast_sort = stable_sort

List.sortList.stable_sortエイリアスになっているようですね.

 天下り的にList.stable_sortの大まかな実装方針を解説すると,これは与えられたリストを前後で2つに分割し,分割されたサブリストに対しソートを行ってからマージするトップダウンな実装になっています.疑似コードにすると以下の通りです.

let rec stable_sort cmp l =
  let (l1, l2) = (* リストlを前後で2つに分割する *) in
  merge cmp (* cmpによって昇順にソートされたリスト2つをソートする *)
    (stable_sort l1) (stable_sort l2)

この実装方針には特に新しい所は無いと思います.何ならC言語マージソートを書いても(再帰をループに直すような工夫をしなければ)こんな実装になりますよね.では,このList.stable_sortがどのように最適化されているのか,3つの点について見ていきましょう.

メモリ確保の削減

 マージソートの実装の本体sortrev_sortは引数を工夫することで,リストを前後で2つに分割する所でコンスセルの確保が起きないようにされています.該当箇所はソースコードの以下の部分ですね(rev_sortの該当箇所は略).

| n, l ->
   let n1 = n asr 1 in
   let n2 = n - n1 in
   let l2 = chop n1 l in
   let s1 = rev_sort n1 l in
   let s2 = rev_sort n2 l2 in
   rev_merge_rev s1 s2 []

 OCamlにはGCがあるのに何故メモリの話が出るのかと不思議に思う方も居るかもしれません. 確かにOCamlにはGCがありますから,確保したメモリは不要になれば自動的に解放され,プログラマはメモリの管理から解放される,と言われています.しかしGCにかかる実行時間は0ではない訳ですから,GCの頻度を下げたいのも事実.GCの頻度を下げるためには余計な(一時的にしか使わない)メモリの確保を避けなければなりません.

 前に述べたようにList.stable_sortは与えられたリストを前後で2つに分割し,分割されたサブリストに対しソートを行ってからマージを行うトップダウンな実装になっています.その与えられたリストを前後で2つに分割する処理ですが,例えばナイーブに書くと以下のようになるでしょう.

let rec chop_bad k l =
  if k = 0 then ([], l)
  else
    match l with
    | x :: l ->
        let (xs, ys) = chop_bad (k - 1) l in
        (x :: xs, ys)
    | _ -> assert false

このchop_bad k lk番目の要素を境にしてリストlを前後2つに分割する関数です1. ところがこのchop_bad k l,普通のMLの実装では再帰呼び出しを行うごとにコンスセルを確保してしまいます2. わざわざ新しく作るリストはと言うと,一番最後にnilが入っている以外lの前半部分と全く同じものですから,このリストを新しく作るのは無駄なように思えてきます.

 そこで,マージソートの実装sortrev_sortの方を,nilではなく長さでリストの終端位置を管理するように書き換えています.言い換えると,sort n llの前からn要素のソート結果を返すようにした訳ですね.こうすれば要素数nn / 23に置き換えるだけでリストlの前半部分を表せますし,リストの後半部分を取ってくる関数chop k lは新たにコンスセルを確保しないため,リストを前後で2つに分割する操作ではコンスセルの確保が起こらないのです.

末尾再帰によるマージの実装

 厳密に言えば実行効率を高めるためではなく,スタックオーバーフローしないための工夫なのですが,アキュムレータを導入することで,末尾再帰でマージが実装されています.ソースコードではstable_sort中のrev_mergeに対応する部分ですね.

let rec rev_merge l1 l2 accu =
  match l1, l2 with
  | [], l2 -> rev_append l2 accu
  | l1, [] -> rev_append l1 accu
  | h1::t1, h2::t2 ->
     if cmp h1 h2 <= 0
      then rev_merge t1 l2 (h1::accu)
      else rev_merge l1 t2 (h2::accu)

 ただ,アキュムレータを使って末尾再帰にしたために,rev_mergeの返すリストは逆順になってしまっています.要するにリストl1l2のマージをmerge l1 l2と書くと,任意のリストl1l2accについてrev_merge l1 l2 acc = List.rev (merge l1 l2) @ accが成り立つ訳ですね.

 任意のリストlについてList.rev (List.rev l) = lは明らかに成り立ちますから,List.rev (rev_merge l1 l2 [])と書いてもリストl1l2のマージを実装できることでしょう.しかし,List.revは与えられたリストの長さに比例した時間計算量を必要としますから,それを避けるために逆順でソートされたリストについての(末尾再帰な)マージrev_merge_rev,及び逆順のソートrev_sortが用いられています.

let rec rev_merge_rev l1 l2 accu =
  match l1, l2 with
  | [], l2 -> rev_append l2 accu
  | l1, [] -> rev_append l1 accu
  | h1::t1, h2::t2 ->
      if cmp h1 h2 > 0
      then rev_merge_rev t1 l2 (h1::accu)
      else rev_merge_rev l1 t2 (h2::ascu)

 任意のリストl1l2accについてrev_merge_rev (List.rev l1) (List.rev l2) acc = merge l1 l2 @ accが成り立ちますし,リストlの前からn要素のソートをsort n lと書くと,任意のリストlについてrev_sort n (List.rev l) = sort n lが成り立ちます.従ってList.rev (List.rev l) = lより,2回の再帰呼び出しで元の順序に戻るわけです.

 ところで,マージの実装rev_merge及びrev_merge_revは末尾再帰になっているのに,マージソートの実装sort及びrev_sortは末尾再帰になっていません.これはsort l及びrev_sort lが高々 O(\log |l|)回の再帰呼び出ししか行わないため,実用上スタックオーバーフローの危険は無いと考えてのことでしょう.

短いリストに対するチューニング

 正直これは泥臭い最適化の部類なのですが,ソートするリストが十分に短くなった場合は再帰呼び出しを止め,ソートの処理をべた書きしています.以下の部分ですね(rev_sortの該当部分は略).

let rec sort n l =
  match n, l with
  | 2, x1 :: x2 :: _ ->
     if cmp x1 x2 <= 0 then [x1; x2] else [x2; x1]
  | 3, x1 :: x2 :: x3 :: _ ->
     if cmp x1 x2 <= 0 then begin
       if cmp x2 x3 <= 0 then [x1; x2; x3]
       else if cmp x1 x3 <= 0 then [x1; x3; x2]
       else [x3; x1; x2]
     end else begin
       if cmp x1 x3 <= 0 then [x2; x1; x3]
       else if cmp x2 x3 <= 0 then [x2; x3; x1]
       else [x3; x2; x1]
     end

まぁでもこういう最適化はクイックソートの実装なんかでもやってるんじゃないでしょうか.

まとめ

 List.sortの実装,いかがでしたでしょうか.普段良く使っている標準ライブラリの関数がこんなにも最適化されているだなんて,僕は思いもしませんでした.あんまり多用すると早すぎる最適化になってしまいそうですが,GCの回数を減らすためにメモリの確保を減らすのは興味深いですね.皆さんもOCaml標準ライブラリのコードを読むと,何か発見が得られるかもしれませんよ?


  1. 厳密に言うと,chop_bad k lが関数なのではなく,chop_badが関数なんですけども.

  2. 解説のための例なので,ここでは2つ組の確保は無視します.

  3. 実際のコードではシフト演算で計算されています.って言うか,普通の言語処理系ならこの程度の最適化はやってくれると思うんですけどね.