Skip to content

Instantly share code, notes, and snippets.

@MikuroXina
Last active November 16, 2023 00:13
Show Gist options
  • Save MikuroXina/1c749ca085b6325cfc11b0373b6c791b to your computer and use it in GitHub Desktop.
Save MikuroXina/1c749ca085b6325cfc11b0373b6c791b to your computer and use it in GitHub Desktop.
Find cumulative sum of numbers with Rust.
pub fn cum<T: Default + std::ops::AddAssign + Copy>(
nums: impl IntoIterator<Item = T>,
) -> impl Iterator<Item = T> {
nums.into_iter().scan(T::default(), |state, item| {
*state += item;
Some(*state)
})
}
pub fn cum2<T: Default + std::ops::AddAssign + Copy, NNS, NS>(nums: NNS) -> Vec<Vec<T>>
where
NS: IntoIterator<Item = T>,
NS::IntoIter: ExactSizeIterator,
NNS: IntoIterator<Item = NS>,
{
nums.into_iter()
.scan(None, |state, item| {
let item_iter = item.into_iter();
if state.is_none() {
*state = Some(vec![T::default(); item_iter.len()]);
}
for (state, cell) in state.as_mut().unwrap().iter_mut().zip(cum(item_iter)) {
*state += cell;
}
state.as_ref().cloned()
})
.collect()
}
#[test]
fn test_cum() {
assert_eq!(
vec![0, 1, 3, 6, 10],
cum(vec![0, 1, 2, 3, 4]).collect::<Vec<_>>()
);
}
#[test]
fn test_cum2() {
assert_eq!(
vec![
vec![0, 0, 1, 1, 2],
vec![0, 1, 2, 4, 5],
vec![0, 1, 5, 7, 8],
vec![0, 0, 4, 7, 8],
],
cum2(vec![
vec![0, 0, 1, 0, 1],
vec![0, 1, 0, 2, 0],
vec![0, 0, 3, 0, 0],
vec![0, -1, 0, 1, 0],
])
);
}
/// Group must satisfy:
/// ```rs
/// fn test<G: Group>(m: G, n: G, l: G) {
/// assert_eq!(m.op(&G::id()), m);
/// assert_eq!(G::id().op(&m), m);
/// assert_eq!(m.op(&n.op(&l)), m.op(&n).op(&l));
/// assert_eq!(m.op(&m.inv()), G::id());
/// assert_eq!(m.inv().op(&m), G::id());
/// }
/// ```
pub trait Group {
fn id() -> Self;
fn op(&self, other: &Self) -> Self;
fn inv(&self) -> Self;
}
pub fn cum_mut<T: Group>(nums: &mut [T]) {
for i in 1..nums.len() {
nums[i] = nums[i - 1].op(&nums[i]);
}
}
pub fn cum2_mut<T: Group>(nums: &mut [Vec<T>]) {
for y in 1..nums.len() {
for x in 0..nums[y].len() {
nums[y][x] = nums[y - 1][x].op(&nums[y][x]);
}
}
for y in 0..nums.len() {
for x in 1..nums[y].len() {
nums[y][x] = nums[y][x - 1].op(&nums[y][x]);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment