[Rust] クイックソートを並列化

jlkiri さんの「Rustで書くソートアルゴリズム」という記事を読んで、そういえば昔書いた並列クイックソートがどこかに転がっていたなと思い出したので公開しておきます. 二年ほど前に並列計算の練習問題として実装したんですが, 今見直すとかなり気になる点があちこちにあったので公開するにあたり修正しました.

なお ウィキペディア に倣ってピボット選択に先頭, 中央, 最後尾の値の中央値を採用しました. そのためにクレートのルートで median という関数を定義しておきます.

src/lib.rs

pub mod serial;
pub mod parallel;

pub(crate) fn median<T:PartialOrd>(a: T, b: T, c:T) -> T {
    if c > b {
        if b > a { b } else {
            if c > a { a } else { c }
        }
    } else {
        if c > a { c } else {
            if b > a { a } else { b }
        }
    }
}

シリアル版

シリアルコードは, 下手に日本語で解説するよりコードを読んだ方が早いです. なお ウィキペディア のコードをベースにしました.

src/serial.rs

pub fn quicksort<T:PartialOrd+Copy>(vec: &mut [T]) {
    if vec.len() > 1 {
        let mut i: usize = 0;
        let mut j: usize = vec.len() - 1;
        let pivot: T = crate::median( vec[0], vec[j/2], vec[j] );
        loop {
            while vec[i] < pivot { i = i + 1; }
            while vec[j] > pivot { j = j.saturating_sub(1); }
            if i >= j { break; }
            vec.swap(i, j);
            i = i + 1;
            j = j.saturating_sub(1);
        }
        quicksort( &mut vec[..i] );
        quicksort( &mut vec[j+1..] );
    }
}

並列版

THRESHOLD という定数を用意して, 一定以上再帰が深くなったら新規スレッドを立てずシリアルコードを呼ぶようにしています. vec の所有権が若干トリッキーですが, Rust のコンパイル時検査のおかげで並列コードでも安心です.

src/parallel.rs

use std::thread;
const THRESHOLD: usize = 3;

pub fn quicksort<T:'static+PartialOrd+Copy+Send>(vec: &mut Vec<T>) {
    qs_loop( vec, 0 );
}

fn qs_loop<T:'static+PartialOrd+Copy+Send>(vec: &mut Vec<T>, depth: usize) {
    if vec.len() > 1 {
        let mut i: usize = 0;
        let mut j: usize = vec.len() - 1;
        let pivot: T = crate::median( vec[0], vec[j/2], vec[j] );
        loop {
            while vec[i] < pivot { i = i + 1; }
            while vec[j] > pivot { j = j.saturating_sub(1); }
            if i >= j { break; }
            vec.swap(i, j);
            i = i + 1;
            j = j.saturating_sub(1);
        }
        if depth < THRESHOLD { // 並列実行
            let mut vec1: Vec<T> = vec.split_off(j+1);
            let handle = thread::spawn(move || {
                qs_loop( &mut vec1, depth+1 );
                vec1
            });
            qs_loop( vec, depth+1 );
            let vec1 = handle.join().unwrap();
            vec.extend_from_slice(&vec1);
        } else { // シリアル実行
            crate::serial::quicksort( &mut vec[..i] );
            crate::serial::quicksort( &mut vec[j+1..] );
        }
    }
}

ライセンス表記

本記事のクイックソートのコードは クイックソート - Wikipedia のC言語コードをベースに Rust で書き直したものです (オリジナルコードの著作権者については 履歴 をご確認ください). もとのコードは CC BY-SA 3.0 のもとで配布されているため, ライセンスに従い本記事のコードも CC BY-SA 3.0 とします.