How to circumvent `take_while` skipping values? - rust

In trying to chain std::iter::Iterator::take_while calls together I'm losing the last values of each call.
Is there a way to chain calls together like this without skipping values?
Code Playground link:
use std::fmt;
#[derive(Clone)]
struct Point {
value: u8,
xe: u8,
xs: u8,
y: u8,
}
impl fmt::Debug for Point {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.value)
}
}
fn main() {
// All values following 5s within its x distance, partitioned by whether it is above or below.
// Sorted by xs (x start) (xe = x end)
#[rustfmt::skip]
let vec:Vec<Point> = vec![
Point { value: 4, xe: 1, xs: 1, y: 2 }, // 4
Point { value: 3, xe: 3, xs: 2, y: 3 }, // 3
Point { value: 5, xe: 7, xs: 4, y: 6 }, // ---- 5 -----
Point { value: 3, xe: 5, xs: 5, y: 4 }, // 3
Point { value: 6, xe: 6, xs: 6, y: 8 }, // 6
Point { value: 2, xe: 8, xs: 8, y: 3 }, // 2
Point { value: 8, xe: 10, xs: 9, y: 2 }, // 8
Point { value: 5, xe: 15, xs: 10, y: 7 }, // ---- 5 -----
Point { value: 2, xe: 12, xs: 11, y: 10 }, // 2
Point { value: 7, xe: 13, xs: 13, y: 9 }, // 7
Point { value: 4, xe: 14, xs: 14, y: 2 } // 4
];
let mut iter = vec.iter();
loop {
let c: Vec<_> = iter
.by_ref()
.take_while(|x| x.value != 5)
.cloned()
.collect();
println!("c: {:.?}", c);
if let Some(var) = iter.next() {
println!("var: {:.?}", var);
let (a, b): (Vec<_>, Vec<_>) = iter
.by_ref()
.take_while(|x| x.xe < var.xe)
.partition(|x| x.y > var.y);
println!("a: {:.?}", a);
println!("b: {:.?}", b);
} else {
break;
}
}
}
Output:
c: [4, 3]
var: 3
a: []
b: []
c: [2, 8]
var: 2
a: []
b: []
c: [4]
It should output:
c: [4, 3]
var: 5
a: [3]
b: [6]
c: [2, 8]
var: 5
a: [2, 7]
b: [4]
Using take_while with std::iter::Iterator::partition seemed a good way to make the code for this relatively clean.
In context the c, a and b values would be passed to functions whose results would be appended to a return value.

Using next_if() and from_fn():
use std::iter::from_fn;
// ...
let mut iter = vec.iter().peekable();
// ...
let c: Vec<_> = from_fn(|| iter.next_if(|x| x.value != 5))
.cloned()
.collect();
// ...
let (a, b): (Vec<_>, Vec<_>) = from_fn(|| iter.next_if(|x| x.xe < var.xe))
.partition(|x| x.y > var.y);
Using peeking_take_while() (better) or take_while_ref() from itertools, just replace the function.

Related

Sharing reference from a mutable method

Rust does not allow borrowing multiple mutable references. I understand that. But I can not find any elegant way to implement a few algorithms. Below is a simplified version of one such algorithm. The Ladder struct hands out slices of ever increasing sequence of numbers, such as, [0], [0, 1], [0, 1, 2] and so on.
struct Ladder {
position: usize,
data: [u8; 10],
}
impl Ladder {
fn get_next(&mut self) -> &[u8] {
self.position += 1;
&(self.data[0..self.position])
}
fn new() -> Ladder {
Ladder {
position: 0,
data: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
}
}
}
I need to call get_next() a couple of times, collect the returned sequences and call a closure that will do something with those sequences.
fn test_ladder(consumer: impl Fn(&[&[u8]])) {
let mut l = Ladder::new();
let mut steps: [&[u8]; 3] = [&[]; 3];
steps[0] = l.get_next();
steps[1] = l.get_next();
steps[2] = l.get_next();
consumer(&steps);
}
fn main() {
test_ladder(|steps| {
for seq in steps {
println!("{:?}", *seq);
}
});
}
It is a non-allocating algorithm. I can not use std::Vec.
What is the idiomatic way to approach problems like this?
The problem here is that you can't keep references to something that you mutate, and .get_next() is allowed to mutate data. What you need to do is separate the data from the mutation. You can do that by only keeping a reference to the original data.
Creating a sequence of elements sounds a lot like an iterator, so here's an example:
struct LadderIter<'a> {
position: usize,
data: &'a [u8],
}
impl<'a> LadderIter<'a> {
fn new(data: &'a [u8]) -> LadderIter<'a> {
LadderIter { position: 0, data }
}
}
impl<'a> Iterator for LadderIter<'a> {
type Item = &'a [u8];
fn next(&mut self) -> Option<Self::Item> {
if self.position == self.data.len() {
None
} else {
self.position += 1;
Some(&self.data[0..self.position])
}
}
}
Which you can then use as an iterator:
for step in LadderIter::new(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) {
println!("{step:?}");
}
[0]
[0, 1]
[0, 1, 2]
[0, 1, 2, 3]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4, 5]
[0, 1, 2, 3, 4, 5, 6]
[0, 1, 2, 3, 4, 5, 6, 7]
[0, 1, 2, 3, 4, 5, 6, 7, 8]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
Or in your specific use-case:
let data = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
let mut ladder = LadderIter::new(&data);
let steps: [&[u8]; 3] = [
ladder.next().unwrap(),
ladder.next().unwrap(),
ladder.next().unwrap(),
];
Another approach is to use interior mutability. Since you are only modifying position, you can use the zero-cost Cell:
use std::cell::Cell;
struct Ladder {
position: Cell<usize>,
data: [u8; 10],
}
impl Ladder {
fn get_next(&self) -> &[u8] {
self.position.set(self.position.get() + 1);
&self.data[0..self.position.get()]
}
fn new() -> Ladder {
Ladder {
position: Cell::new(0),
data: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
}
}
}

Cannot borrow `counts` as immutable because it is also borrowed as mutable

I'm getting this error
error[E0502]: cannot borrow `counts` as immutable because it is also borrowed as mutable
--> src/main.rs:13:64
|
13 | counts.entry(x).or_insert(cmp::max(current_count, *counts.get(&x).unwrap()));
| --------------- --------- ^^^^^^^^^^^^^^ immutable borrow occurs here
| | |
| | mutable borrow later used by call
| mutable borrow occurs here
but I can't seem to understand what's wrong despite having seen a few other questions asking about the same error and they are using slice instead of dealing with a HashMap like here.
The 2nd argument of cmp::max has to be a value, but I cant figure out how to make use of an mutable borrow here as suggested by the compiler.
use std::collections::HashMap;
use std::cmp;
fn main() {
let data: Vec<u8> = vec![1,1,3,1,1];
let mut counts = HashMap::new();
let mut current_count = 0;
for (i, &x) in data.iter().enumerate() {
if i > 0 {
if x == data[i-1] {
current_count += 1;
counts.entry(x).or_insert(cmp::max(current_count, *counts.get(&x).unwrap()));
} else {
current_count = 0;
}
}
}
println!("{:?}", counts);
}
#BallpointBen I think this might work, but not sure if this is a good way.
for (i, &x) in data.iter().enumerate() {
if i > 0 {
if x == data[i-1] {
current_count += 1;
counts.entry(x).or_insert(1);
counts.insert(x, cmp::max(current_count, counts[&x]));
} else {
current_count = 0;
}
}
}
HashMap::entry() borrows the map mutably for the whole life of the returned Entry. You cannot use the map until you finish with the entry.
The fix is simple - first retrieve the value into a variable, then call entry():
let v = *counts.get(&x).unwrap();
counts
.entry(x)
.or_insert(cmp::max(current_count, v));
You cannot use or_insert_with() by that, which would be preferred, but you didn't use it anyway.
However, what you're trying to do is not good: when will your or_insert() be needed? Only when there is no value associated with this key in the map, so trying to get().unwrap() it is always going to panic!
Here is a solution that meshes well with the type system. We append a None to the end of the iterator to signal that we're done (roughly equivalent to an if i == my_vec.len()-1 check but without needing a Vec)
use std::collections::HashMap;
use std::hash::Hash;
fn longest_run<T: Hash + Eq>(items: impl IntoIterator<Item = T>) -> HashMap<T, usize> {
let mut runs = HashMap::new();
let mut items = items.into_iter();
let mut prev = match items.next() {
Some(x) => x,
None => return runs,
};
let mut curr_count = 1;
for item in items.map(Some).chain(std::iter::once(None)) {
if let Some(item) = item {
if item == prev {
curr_count += 1;
} else {
let prev_count = runs.entry(prev).or_insert(1);
*prev_count = (*prev_count).max(curr_count);
curr_count = 1;
}
prev = item;
} else {
let prev_count = runs.entry(prev).or_insert(1);
*prev_count = (*prev_count).max(curr_count);
// Redundant, but needed to tell the compiler it's
// ok to move prev into `entry` without
// reinitializing it, as we won't need it again
break;
}
}
runs
}
fn main() {
for v in vec![
vec![],
vec![1],
vec![1, 1, 1, 1, 1],
vec![1, 2],
vec![1, 1, 2, 1],
vec![1, 1, 2, 1, 1, 1, 1, 2, 2, 2],
vec![1, 1, 1, 1, 1, 2, 3, 1, 1, 2, 2, 3, 3],
] {
println!("{:?}: {:?}", v, longest_run(v.iter().copied()));
}
}
Prints
[]: {}
[1]: {1: 1}
[1, 1, 1, 1, 1]: {1: 5}
[1, 2]: {2: 1, 1: 1}
[1, 1, 2, 1]: {1: 2, 2: 1}
[1, 1, 2, 1, 1, 1, 1, 2, 2, 2]: {2: 3, 1: 4}
[1, 1, 1, 1, 1, 2, 3, 1, 1, 2, 2, 3, 3]: {2: 2, 3: 2, 1: 5}

Permutations with replacement in rust?

I'd like to write a generic function that can make all unique vectors where each index has a series of values.
Easiest to illustrate with an example.
for i in (1..3).combinations_with_replacement(3) {
println!("{:?}",i);
}
Produces the ouput
[1, 1, 1]
[1, 1, 2]
[1, 2, 2]
[2, 2, 2]
Which is not satisfactory because it's missing members like
[1, 2, 1]
[2, 1, 2]
[2, 2, 1]
So I also tried permutations(3) but since their are more positions than items in the iterator, the iterator is empty. There also doesn't appear to be a permutations_with_replacement but maybe that'd be the name of the function I'm looking for.
In this case you could accomplish the task with 3 nested for loops but this is ugly and not a general solution. I think a recursive backtracking solution could do it too but seems like their should be something in itertools that I'm missing.
Here's another example of what I want but with code written a few languages that aren't rust.
Using itertools, based on #Ry-'s suggestion:
for i in (1..=3).map(|_| 1..=2).multi_cartesian_product() {
println!("{i:?}");
}
Output:
[1, 1, 1]
[1, 1, 2]
[1, 2, 1]
[1, 2, 2]
[2, 1, 1]
[2, 1, 2]
[2, 2, 1]
[2, 2, 2]
Playground
I'm not sure how this can be done using itertools, but here is one way to achieve the same result in plain Rust.
pub struct PermutationsReplacementIter<I> {
items: Vec<I>,
permutation: Vec<usize>,
group_len: usize,
finished: bool,
}
impl<I: Copy> PermutationsReplacementIter<I> {
fn increment_permutation(&mut self) -> bool {
let mut idx = 0;
loop {
if idx >= self.permutation.len() {
return true;
}
self.permutation[idx] += 1;
if self.permutation[idx] >= self.items.len() {
self.permutation[idx] = 0;
idx += 1;
} else {
return false;
}
}
}
fn build_vec(&self) -> Vec<I> {
let mut vec = Vec::with_capacity(self.group_len);
for idx in &self.permutation {
vec.push(self.items[*idx]);
}
vec
}
}
impl<I: Copy> Iterator for PermutationsReplacementIter<I> {
type Item = Vec<I>;
fn next(&mut self) -> Option<Self::Item> {
if self.finished {
return None;
}
let item = self.build_vec();
if self.increment_permutation() {
self.finished = true;
}
Some(item)
}
}
pub trait ToPermutationsWithReplacement {
type Iter;
fn permutations_with_replacement(self, group_len: usize) -> Self::Iter;
}
impl<I: Iterator> ToPermutationsWithReplacement for I {
type Iter = PermutationsReplacementIter<<I as Iterator>::Item>;
fn permutations_with_replacement(self, group_len: usize) -> Self::Iter {
let items = self.collect::<Vec<_>>();
PermutationsReplacementIter {
permutation: vec![0; group_len],
group_len,
finished: group_len == 0 || items.len() == 0,
items,
}
}
}
Then it can be used similarly to combinations_with_replacement.
for x in (1..3).permutations_with_replacement(3) {
println!("{:?}", x);
}
// Output:
[1, 1, 1]
[2, 1, 1]
[1, 2, 1]
[2, 2, 1]
[1, 1, 2]
[2, 1, 2]
[1, 2, 2]
[2, 2, 2]
Playground Link
You can also put it on any iterator where the elements implement Copy. However, I wouldn't recommend it. The time complexity for this task is extremely bad. For an input iterator of n elements in groups of length m, this will create an iterator of about n^m items assuming my math is correct.

Insert into Rust array in place, push other elements down

I'm trying to do the following in Rust, specifically using arrays (I don't want to use vectors here, and want elements pushed out of the array if we're done).
let mut x = [1, 2, 3, 4, 5];
// array, number to insert, place to be inserted at
insert_in_place(&x, 7, 1);
// x is now [1, 7, 2, 3, 4];
How do you implement insert_in_place?
I think there's a way to do this using slices, but I'm still learning and wondering if there's a really elegant way to do this kind of thing.
fn insert_in_place<T>(array: &mut [T], value: T, index: usize) {
*array.last_mut().unwrap() = value;
array[index..].rotate_right(1);
}
Try it online!
Or equivalently:
fn insert_in_place<T>(array: &mut [T], value: T, index: usize) {
array[index..].rotate_right(1);
array[index] = value;
}
Try it online!
Iterate the slice, skipping elements before the index of the one you need to insert. Then swap each element with its previous element (or, for the first one, use the item to add).
fn insert_in_place<T>(x: &mut [T], new: T, index: usize) {
let mut next = new;
for e in x.iter_mut().skip(index) {
std::mem::swap(e, &mut next);
}
}
fn main() {
let mut x = [1, 2, 3, 4, 5];
// array, number to insert, place to be inserted at
insert_in_place(&mut x, 7, 1);
// x is now [1, 7, 2, 3, 4];
println!("{:?}", x);
}

Cycle a Rust iterator a given number of times

How do I cycle through an iterator a finite number of times?
I would expect the output of something like this to be 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3 and then stop:
vec![1, 2, 3].iter().cycle(4)
// ^ but .cycle() doesn't take an argument...
I don't know the length of the iterator to begin with.
One simple way is to repeat the iterator itself, take the first 4 and flatten:
fn main() {
let v = vec![1, 2, 3];
let res = std::iter::repeat(v.iter())
.take(4)
.flatten()
.collect::<Vec<_>>();
dbg!(res);
}
Some micro-benchmark result using code in this gist comparing 3 different approaches:
repeat-take-flatten in this answer
hand-rolled loop
a cycle_n implementation mimicking Iterator::cycle.
Kudos to rustc, cycle_n consistently outperforms the other two when the input is reasonably large whereas repeat-take-flatten performs the best for small input.
There is no such an iterator in the std lib.
If you know the iterator size, you can take your number times the size of the iterator:
fn cycle_n_times<T: Clone>(slice: &[T], count: usize) -> impl Iterator<Item = &T> {
slice.iter().cycle().take(slice.len() * count)
}
Or you can write your own that is more general:
pub struct Ncycles<I> {
orig: I,
iter: I,
count: usize,
}
impl<I: Clone> Ncycles<I> {
fn new(iter: I, count: usize) -> Ncycles<I> {
Ncycles {
orig: iter.clone(),
iter,
count,
}
}
}
impl<I> Iterator for Ncycles<I>
where
I: Clone + Iterator,
{
type Item = <I as Iterator>::Item;
#[inline]
fn next(&mut self) -> Option<<I as Iterator>::Item> {
match self.iter.next() {
None if self.count == 0 => None,
None => {
self.iter = self.orig.clone();
self.count -= 1;
self.iter.next()
}
y => y,
}
}
}
#[test]
fn it_work() {
Ncycles::new(vec![1, 2, 3].iter(), 4).eq(&[1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3]);
}
One could leverage slice::repeat, but I can't imagine this is very efficient:
let v = vec![1, 2, 3];
let it = v.iter();
println!("{:?}", it.map(|x| *x).collect::<Vec<i32>>().repeat(4).iter());

Resources