Constrained Iterator Sum

1

I've lately been hitting cases where I have an Iterator<Item = usize>, and want to write the equivalent of it.sum() == 2 or it.sum() < 4.

Now, obviously I could just write that, but it seems silly to potentially sum the whole of the iterator, which could be many many items, if the first element is already 10.

Is there a tidy way to do that with iterator adapters? I couldn't really come up with anything better than something like .try_fold(0, |mut sum, x| { sum += x; if sum == 2 { Break(true) } else if sum > 2 { Break(false) } else { Continue(sum) }) which just doesn't feel elegant.

Answers also welcome for the simpler cases of .count() == 2 or .count() < 4 -- that enables other possibilities like .enumerate().any(|(i, _)| i+1 == 2), thought that still feels pretty obtuse.

2

Well, I’d start by using Iterator::scan() to define running_total(), which makes sum() < 4 easy:

use core::ops;
use num_traits::Zero;

fn running_total<T, I>(i: I) -> impl Iterator<Item = T>
where
    T: ops::AddAssign + Copy + Zero,
    I: IntoIterator<Item = T>,
{
    i.into_iter().scan(T::zero(), |sum, val| {
        *sum += val;
        Some(*sum)
    })
}

fn sum_less_than(limit: usize, i: impl IntoIterator<Item = usize>) -> bool {
    running_total(i).all(|sum| sum < limit)
}

fn main() {
    assert_eq!(running_total([1, 2, 3]).collect::<Vec<_>>(), vec![1, 3, 6]);
    assert!(sum_less_than(7, [1, 2, 3]));
    assert!(!sum_less_than(6, [1, 2, 3]));
    assert!(!sum_less_than(3, [1, 2, 3]));
}

I don’t have a really great idea for equality, though — I think it kind of has to be three cases (continue, break true, break false). But if you build on running_total(), you don't need it to be stateful, at least.

4 Likes
3

Thanks, I really like that rephrasing of it. I wasn't thinking of all, but of course if I was going to write a loop manually for this I'd be checking the threshold on every item, so all makes perfect sense.

And then for my .count() < 4 example, I realized that there's already a function that's perfect for that: .nth(4).is_none().

4

Note that short-circuiting a sum only makes sense for unsigned types, to make sure the cumulative sum is monotonically increasing.

For sum equality, Itertools::find_or_last helps:

fn sum_equal(target: usize, i: impl IntoIterator<Item = usize>) -> bool {
    running_total(i)
        .find_or_last(|&sum| sum > target)
        .unwrap_or(0)
        == target
}