Rustでk-permutations

最近はRustの勉強がてらAtcoderやAOJの簡単な問題を解いているが、その際に順列を生成する関数が欲しくなったので実装してみた (Atcoderでは itertools クレートを使えるので不要だが)。

実装全体はこちら: Generate k-permutations in Rust

動き方としては、

let s = "ABCD";
s.chars().permutations(2).for_each(|cs| {
    println!("{:?}", cs);
});
// Output:
// ['A', 'B']
// ['A', 'C']
// ['A', 'D']
// ['B', 'A']
// ['B', 'C']
// ['B', 'D']
// ['C', 'A']
// ['C', 'B']
// ['C', 'D']
// ['D', 'A']
// ['D', 'B']
// ['D', 'C']

のようにIteratorから permutations(k) を呼ぶことで、Iteratorからk個取り出して作る順列を (indexオーダーでの) 辞書順で順番に返すようになっている。

実装は、

struct PermutationIterator<T> {
    v: Vec<(usize, T)>, // next()内でv[..k]が次の順列となるように要素をスワップされるベクトル。(index, 要素)の2つ組
    k: usize,           // nPkのk
    finished: bool,     // 順列生成が終了したかどうかのフラグ
}

のような構造体を作り、以下のようにこの構造体に next() を実装してIteratorとしての振る舞いを定義している。

impl<T: Copy> Iterator for PermutationIterator<T> {
    type Item = Vec<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            None
        } else {
            let res = self.v.iter().map(|e| e.1).take(self.k).collect::<Vec<T>>();
            self.next_k_permutation();
            Some(res)
        }
    }
}

next() を呼ぶたびにIteratorから返却されるベクトルを

let res = self.v.iter().map(|e| e.1).take(self.k).collect::<Vec<T>>();

で作ってから self.next_k_permutation() で v[..k] が次の辞書順になるように変換する。繰り返し順列を生成し、 v[..k] の辞書順が生成される順列内で最大になったときに finished = true にしてIteratorが終了するようにしている。

v[..k] を辞書順で生成する next_k_permutation() は Stack Overflowの回答を参考にして書いた。

impl<T> PermutationIterator<T> {
    fn flip(&mut self, mut lo: usize, mut hi: usize) {
        while lo + 1 < hi {
            self.v.swap(lo, hi - 1);
            lo += 1;
            hi -= 1;
        }
    }

    fn next_k_permutation(&mut self) {
        let mut tail_max = self.v[self.n() - 1].0;
        let mut tail_start = self.k;
        while tail_start > 0 && self.v[tail_start - 1].0 >= tail_max {
            tail_start -= 1;
            tail_max = self.v[tail_start].0;
        }
        if tail_start == 0 {
            self.finished = true;
        } else {
            let mut swap_in: usize;
            let pivot = self.v[tail_start - 1].0;
            if pivot >= self.v[self.n() - 1].0 {
                swap_in = tail_start;
                while swap_in + 1 < self.k && self.v[swap_in + 1].0 > pivot {
                    swap_in += 1;
                }
            } else {
                swap_in = self.n() - 1;
                while swap_in > self.k && self.v[swap_in - 1].0 > pivot {
                    swap_in -= 1;
                }
            }
            self.v.swap(tail_start - 1, swap_in);
            self.flip(self.k, self.n());
            self.flip(tail_start, self.n());
        }
    }

    fn n(&self) -> usize {
        self.v.len()
    }
}

self.v を v[..k] と v[k..] に分けて考え、 v[k..] が常に昇順に整列されるように要素を交換しながら v[..k] で辞書順で次のVecを表現している。動作が分かり辛いがコードを追いながら紙にv がどう変わっていくかを書いていけば分かると思う。

たとえば (0..5).permutations(2).for_each(...) としたときの PermutationIterator の v の変化を順に書くとこんな感じである↓

↓順列    ↓常に昇順になっている
v[..2] | v[2..]
1    | 2 3 4
2    | 1 3 4
3    | 1 2 4
4    | 1 2 3
0    | 2 3 4
2    | 0 3 4
3    | 0 2 4
4    | 0 2 3
0    | 1 3 4
1    | 0 3 4
3    | 0 1 4
4    | 0 1 3
0    | 1 2 4
1    | 0 2 4
2    | 0 1 4
4    | 0 1 2
0    | 1 2 3
1    | 0 2 3
2    | 0 1 3
3    | 0 1 2

あとはIteratorから .permutations() を呼べるようにすれば完成だが、これは

trait Permutations<T> {
    fn permutations(self, size: usize) -> PermutationIterator<T>;
}

impl<T, I: IntoIterator<Item = T>> Permutations<T> for I {
    fn permutations(self, k: usize) -> PermutationIterator<T> {
        PermutationIterator::new(self.into_iter(), k)
    }
}

のように permutations() のメソッドシグネチャを持つトレイトを定義し、Iteratorを実装するジェネリクスに対してそのトレイトを実装するコードを書けば実現できる。

大雑把に仕組みをまとめると、Iteratorから .permutations() を呼ぶと PermutationIterator というトレイトが返るようになっており、このトレイトはIteratorとしての振る舞いを持っていて next() を呼ぶたびに内部で next_k_permutation() が呼ばれて辞書順に順列を返す、というものである。

Iteratorから permutations のメソッドを生やす形にしたのと、IteratorのItemに対する制限がCopyトレイトの実装のみなので割と使い勝手が良さそうかな、と思っている。

参考

実装に関してこちらも参考にさせていただいた。

Rust で順列生成する Iterator の実装

PREVIOUSRustで線形回帰