O(n)時間でソートが完了するバケットソートをHaskellで実装する (2)

はじめに

前回までの実装では、つぎのような制限があった。

「(バケットソートなので比較しないけど)比較対象になる値」以外の値をもとのリストに含めることができない。
- つまり、[(3, "foo"), (7, "bar"), (3, "buz")] のような値のリストを数値のみを比較対象にしてソートすることができない

この記事では、
このような制限をなくして、本当のバケットソートを実装する。

アルゴリズム

前回までで入手した「値 - 出力回数」の情報を使って、あらかじめ「バケツの大きさ」を決めておき、その大きさのバケツに要素をほおりこんでいけばいい。つまり、1回目のスキャンでは「出力回数」を入手して、2回目のスキャンで実際のソートを行うというような2パスのアルゴリズムになる。

ソートする値から「比較する値」を取り出す
「比較する値 - 出力回数」のような配列を作成する (前回実装)
それぞれの大きさのバケツを用意する
- 実際には要素数の配列を用意して
- 「比較する値」のそれぞれに対して「置く場所(index)」を指定する
バケツにそれぞれの要素を入れていく
- 配列の「比較する値」によって指定された位置に要素を配置する
- 配列の「比較する値」によって指定された位置をひとつ進める

(アルゴリズムを言葉で説明するのって、なかなかむずかしいな)

MArrayTools

MArrayTools.hsに関数scanMArrayを追加する。

MArrayTools.hs

scanMArray :: (MArray a e m, Ix i) => (s -> e -> s) -> a i e -> s -> m ([s], s)
scanMArray op a s0 = sma s0 . range =<< getBounds a
        where
        sma s [] = return ([], s)
        sma s (i : is) = do
                e <- readArray a i
                first (s :) <$> sma (s `op` e) is

scanlと類似した関数。配列のなかの要素をひとつずつ調べながら、状態をあらわす値を変化さてている。その「変化していく状態」の含化のそれぞれの段階をリストにまとめたものを返す。

最後の状態だけをリストではなくタプルの第2要素として返しているのは、あとで使いやすいようにというだけの話。

BucketSortM

BucketSortM.hsにつぎのような関数bucketSortStep2を追加する。

BucketSortM.hs

bucketSortStep2 :: forall a a' m i x . (MArray a Int m, MArray a' x m, Ix i) =>
        (x -> i) -> [x] -> a i Int -> m (a' Int x)
bucketSortStep2 getIx lst ns = do
        bs <- getBounds ns
        (jsgen, len) <- scanMArray (+) ns 0
        js <- newListArray bs jsgen :: m (a i Int)
        xs <- newArray_ (0, len - 1)
        for_ lst $ \x -> do
                let   i = getIx x
                j <- readArray js i
                modifyArray js i (+ 1)
                writeArray xs j x
        return xs

配列nsは「比較対象となる値 - 出力回数」のような配列だ。これからscanMArrayによって、配列上の位置のリストjsgenと「ソート対象である要素」の数lenを取り出す。たとえば「0 - 1, 1 - 3, 2 - 1, 3 - 4」のような配列からは、つぎのような結果になる。

jsgen: [0, 1, 4, 5]
len: 9

リストjsgenの意味としては、つぎようになる。

「比較される値0に対応する値は「配列上のインデックス0」から置いていけばいい。おなじように値1に対応する値はインデックス1から、値2に対応する値はインデックス4から、値3に対応する値はインデックス5から置いていけばいい。」

jsgenからnewListArrayで「比較対象である値 - 結果の配列のインデックス」のような配列jsを作成している。つぎに使用されている関数newArray_は「不定の値」を初期値としる配列を作成する。ここで作成されているのは結果を格納する配列xsだ。

関数for_は第1引数であるリスト(本当はリストに限らないけど)のそれぞれの要素に対して、第2引数であるモナド(本当はモナドに限らないけど)を「実行」する。モナドのなかでは、つぎのような処理が行われる。

値xに対する「比較対象の値」であるiを取り出す
配列jsから値iに対応する「置く場所」である位置jを取り出す
配列jsの値iに対応する「置く場所」を、1進める
位置jに値xを置く

最後に「ソートされた配列」である配列xsを返す。

前回定義した関数bucketSortMArrayと、このbucketSortStep2とをつなげて、関数bsortMを定義する。

BucketSortM.hs

bsortM :: forall a a' m i x . (MArray a Int m, MArray a' x m, Ix i) =>
        (x -> i) -> (i, i) -> [x] -> m [x]
bsortM getIx bs xs = getElems
        =<< bucketSortStep2 @a @a' gtIx xs
        =<< bucketSortMArray bs (getIx <$> xs)

bucketSortMArrayで作成した「値の出力回数」を記録した配列をbucketSortStep2にわたして、そこで作られた配列からgetElemsで値を取り出している。

BucketSort

IOモナドで実行できるbsortIOと、純粋な関数bsortとを定義する。だいたい前回のbucketSortIOとbucketSortとおなじ感じ。

BucketSort.hs

bsort :: Ix i => (x -> i) -> (i, i) -> [x] -> [x]
bsort getIx bs xs = runST $ bsortST getIx bs xs

bsortST :: forall s i x . Ix i => (x -> i) -> (i, i) -> [x] -> ST s [x]
bsortST = bsortM @(STArray s) @(STArray s)

bsortIO :: Ix i => (x -> i) -> (i, i) -> [x] -> IO [x]
bsortIO = bsortM @IOArray @IOArray

まとめ

前回実装した「それぞれの値の出力回数」を記録した配列をもとめる関数を使って、それぞれの値に対応する大きさのバケツを用意して、そこに要素をつめこんでいった。

「まちがい」や「不備」、あるいは「こうしたほうがいい」というアドバイスは大歓迎です。ただし、「優しく」指摘していただけると幸いです。

また、質問は大歓迎です。この記事は「突貫」で作ったので、だいぶ説明を簡素化していますし、説明の質自体もそれほど高くはないです。今後、時間が許すかぎりにおいて、洗練していこうかと思いますが、現時点で「ここがわからない」などありましたら、ご質問ください。できる範囲で返答していきたいと思います。