Last active
May 1, 2022 17:54
-
-
Save fhiyo/90da9c9b64b279213a729099497c4574 to your computer and use it in GitHub Desktop.
Generate k-permutations in Rust
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// rustc 1.48.0 (7eac88abb 2020-11-16) | |
struct PermutationIterator<T> { | |
v: Vec<(usize, T)>, | |
k: usize, | |
finished: bool, | |
} | |
impl<T> PermutationIterator<T> { | |
fn new<I: Iterator<Item = T>>(iter: I, k: usize) -> PermutationIterator<T> { | |
let v: Vec<(usize, T)> = iter.enumerate().collect(); | |
if k > v.len() { | |
panic!( | |
"k must be smaller than vec size. k: {}, vec size: {}", | |
k, | |
v.len() | |
); | |
} | |
let finished = v.is_empty() || k == 0; | |
PermutationIterator { v, k, finished } | |
} | |
fn flip(&mut self, mut lo: usize, mut hi: usize) { | |
while lo + 1 < hi { | |
self.v.swap(lo, hi - 1); | |
lo += 1; | |
hi -= 1; | |
} | |
} | |
// ref: https://stackoverflow.com/questions/51287171/generating-k-permutations-from-n-in-c#answer-51292710 | |
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() | |
} | |
} | |
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>>(); | |
// let res = self.v.iter().map(|e| e.1).collect::<Vec<T>>(); | |
self.next_k_permutation(); | |
Some(res) | |
} | |
} | |
} | |
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) | |
} | |
} | |
#[cfg(test)] | |
mod tests { | |
use crate::Permutations; | |
#[test] | |
fn n_1_k_0() { | |
let actual: Vec<Vec<i32>> = (0..1).permutations(0).collect(); | |
let expected: Vec<Vec<i32>> = vec![]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
fn n_1_k_1() { | |
let actual: Vec<Vec<i32>> = (0..1).permutations(1).collect(); | |
let expected = vec![vec![0]]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
#[should_panic] | |
fn n_1_k_2() { | |
let _actual: Vec<Vec<i32>> = (0..1).permutations(2).collect(); | |
} | |
#[test] | |
fn n_3_k_1() { | |
let actual: Vec<Vec<i32>> = (0..3).permutations(1).collect(); | |
let expected = vec![vec![0], vec![1], vec![2]]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
fn n_3_k_2() { | |
let actual: Vec<Vec<i32>> = (0..3).permutations(2).collect(); | |
let expected = vec![ | |
vec![0, 1], | |
vec![0, 2], | |
vec![1, 0], | |
vec![1, 2], | |
vec![2, 0], | |
vec![2, 1], | |
]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
fn n_3_k_3() { | |
let actual: Vec<Vec<i32>> = (0..3).permutations(3).collect(); | |
let expected = vec![ | |
vec![0, 1, 2], | |
vec![0, 2, 1], | |
vec![1, 0, 2], | |
vec![1, 2, 0], | |
vec![2, 0, 1], | |
vec![2, 1, 0], | |
]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
fn n_5_k_2_chars() { | |
let actual: Vec<Vec<char>> = "hello".chars().permutations(2).collect(); | |
let expected = vec![ | |
vec!['h', 'e'], | |
vec!['h', 'l'], | |
vec!['h', 'l'], | |
vec!['h', 'o'], | |
vec!['e', 'h'], | |
vec!['e', 'l'], | |
vec!['e', 'l'], | |
vec!['e', 'o'], | |
vec!['l', 'h'], | |
vec!['l', 'e'], | |
vec!['l', 'l'], | |
vec!['l', 'o'], | |
vec!['l', 'h'], | |
vec!['l', 'e'], | |
vec!['l', 'l'], | |
vec!['l', 'o'], | |
vec!['o', 'h'], | |
vec!['o', 'e'], | |
vec!['o', 'l'], | |
vec!['o', 'l'], | |
]; | |
assert_eq!(expected, actual); | |
} | |
#[test] | |
fn n_3_k_2_mystruct() { | |
#[derive(Debug, Copy, Clone, PartialEq)] | |
struct MyStruct(i32); | |
let actual: Vec<Vec<MyStruct>> = vec![MyStruct(0), MyStruct(1), MyStruct(2)] | |
.permutations(2) | |
.collect(); | |
let expected = vec![ | |
vec![MyStruct(0), MyStruct(1)], | |
vec![MyStruct(0), MyStruct(2)], | |
vec![MyStruct(1), MyStruct(0)], | |
vec![MyStruct(1), MyStruct(2)], | |
vec![MyStruct(2), MyStruct(0)], | |
vec![MyStruct(2), MyStruct(1)], | |
]; | |
assert_eq!(expected, actual); | |
} | |
} | |
fn main() { | |
let chars = "hello".chars(); | |
chars.permutations(3).for_each(|cs| { | |
println!("{:?}", cs); | |
}) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment