2021/01/26

リストを2つのグループに分ける関数を作る

partition

OCamlのListモジュールにはpartitionという便利な関数がありまして、これを使うとクイックソートが簡単に書くことができます。


let rec qsort = function
  | [] -> []
  | h :: t ->
    let (l, r) = List.partition (fun x -> x < h) t in
    (qsort l)@[h]@(qsort r)

こりゃ便利!と思い、他の言語でも使いたい!ということで自前で書くことにしました。

Common Lisp編

まずはCommon Lispで書いてみました。以下はpartition関数とクイックソートのコードです。


(defun %partition (p l)
  (labels ((aux (p y n l)
             (if (endp l)
               (list (nreverse y) (nreverse n))
               (if (funcall p (car l))
                 (aux p (cons (car l) y) n (cdr l))
                 (aux p y (cons (car l) n) (cdr l))))))
    (let ((result (aux p '() '() l)))
      result)))
(defun qsort2 (lst)
  "Quick Sort version 2"
  (cond
    ((endp lst) lst)
    (t
      (let* ((result (%partition #'(lambda (x) (< x (car lst))) (cdr lst)))
             (l (car result))
             (r (cadr result)))
        (nconc (qsort2 l) (cons (car lst) (qsort2 r)))))))

Version 2となっていますのは、もっとすんなりかけたクイックソートがあるからです。詳しいコードはGitHubに置いてます。

Python編

Pythonではリスト内包表記ができるので、以下のような見やすい感じになりました。


def partition(p, l):
    if l == []:
        return l
    else:
        y = [i for i in l if p(i)]
        n = [i for i in l if not p(i)]
        return (y, n)

が、これをごにょごにょっと改造するとクイックソートになります。


def quick_sort(l):
    if l == []:
        return l
    else:
        pivot = l[len(l) // 2]
        y = quick_sort([i for i in l if i < pivot])
        p = [i for i in l if i == pivot]
        n = quick_sort([i for i in l if i > pivot])
        y += p
        y += n
        return y

クイックソートの肝であるピボットは、中央値が良いと風の噂で聴いたのでそのようにしています。

Rust編

Rustでは以下のようになりました。


fn partition<P, T>(predicate: P, list: &mut [T])
    -> Option<(Vec<&T>, Vec<&T>)>
    where P: Fn(T) -> bool,
          T: Copy
{
    match *list {
        [] => None,
        _ => {
            let l = list.iter()
                .filter(|x| predicate(**x))
                .collect::<Vec<&T>>();
            let r = list.iter()
                .filter(|x| !predicate(**x))
                .collect::<Vec<&T>>();
            Some((l.to_vec(), r.to_vec()))
        }
    }
}

が、何となくごちゃごちゃしている感じがしているので、もうちょっと考えて作ってみようと思っています。現在のものでも以下のように、問題なく動作します。


fn main() {
    let mut v = vec![1,2,3,4,5,6,7];
    println!("{:?}", partition(|x| x % 2 == 0, &mut v));

    let mut alpha = vec!["abc", "abd", "bce", "ace"];
    println!("{:?}", partition(|x| x.starts_with("a"), &mut alpha));
}


=======
出力結果
=======
Some(([2, 4, 6], [1, 3, 5, 7]))
Some((["abc", "abd", "ace"], ["bce"]))

しかし、Rustの場合、これを使ってクイックソートをするというのは非常にややこしい。Rustでクイックソートも何だかかなりややこしい。と勉強不足の私はぼやきながら、色々考え中です。すっきり解決したらまた記事にしようと思っています。

追記(2020/01/27)

Rustにはpartition関数がありました。調べたところが違っていてiterのところにありました。


fn main() {
    let v = vec![1,2,3,4,5,6];
    let (l, r): (Vec<i32>, Vec<i32>) = v.iter().partition(|x| *x % 2 == 0);
    println!("{:?}, {:?}", l, r);
}

また無駄な時間を使ってしまつた……