cumsum with reset points

cumsum_cut.R

partition <- function(x, cut) {
  v <- rep(0, length(x))
  v[cut] <- 1
  cumsum(v)
}

cumsum_cut <- function(x, cut = NULL) {
    cut <- sort(cut)
    vs <- split(x, partition(x, cut))
    vs <- sapply(vs, \(x) {x[1] <- 0; x})
    unlist(sapply(vs, cumsum), use.names = FALSE) 
}

cumsum_cut(c(1,1,1,1,1,1,1,1), cut = c(4, 7))
#> [1] 0 1 2 0 1 2 0 1
#> c(0, 1, 2, 0, 1, 2, 0, 1)

set.seed(1)
vec <- sample(80)
reset_at <- sample(80, 3)
res <- cumsum_cut(vec, reset_at)

res
#>  [1]    0   39   40   74  117  131  190  241  262  316  323  360  429  506  550
#> [16]  583  603  638  644  654  696  734  781  845  873  930  996 1019 1081 1121
#> [31] 1192 1217 1269    0   60  108  132  164  239  241  259  281  327  405  450
#> [46]  503  539  550  567  571    0    8   57   73  129  192  266  321  391    0
#> [61]   12   39  119  191  196  211  278  351  360  418  483  486  536  577  638
#> [76]  664  695  724  743  819
res[1]
#> [1] 0
res[reset_at]
#> [1] 0 0 0

# Created on 2024-10-29 with [reprex v2.1.1](https://reprex.tidyverse.org)

cumsum_cut.rs

use std::time::Instant;

fn cumsum_cut(vec: Vec<i64>, cut: Vec<i64>) -> Vec<i64> {
    let mut res: Vec<i64> = Vec::with_capacity(vec.len()+1);
    res.push(0);
    for (i, el) in vec.iter().enumerate() {
        if i == 0 {
            res.push(*el);
        } else if cut.contains(&((i+2) as i64)) {
            res.push(0);
        } else {
            res.push(res[i] + *el);
        }
    }
    return res[0..vec.len()].to_vec()
}

fn main() {
    let start = Instant::now();
    
    println!("{:?}", cumsum_cut(vec!(8, 7, 6, 5, 4, 3, 2, 1), vec!(4, 7)));
 // println!("{:?}", cumsum_cut(vec!(1, 1, 1, 1, 1, 1, 1, 1), vec!(4, 7)));
    
    let elapsed = start.elapsed();
    println!("Time elapsed is: {:?}", elapsed);
}

# cargo run --release
#    Finished `release` profile [optimized] target(s) in 0.00s
#     Running `target/release/cumsum_cut`
# [0, 8, 15, 0, 5, 9, 0, 2]
# Time elapsed is: 62.75µs

Here's my current rust implementation

fn fx(v: &mut [i32], cuts: Vec<i32>) -> Vec<i32> {

    // return early 
    if v.len() < 2 {
        return v.to_vec();
    }

    let len = v.len(); 

    // shift elements
    v.copy_within(0..len-1, 1);
    
    // Step 2: Set the first element to 0.
    v[0] = 0;

    // iterate and set cuts to 0
    for index in cuts {
        if index < v.len() as i32 {
            v[(index -1) as usize] = 0;
        }
    }

    // do cumsum
    for i in 1..len {
        if v[i] == 0 {
            continue;
        }
        v[i] += v[i - 1];
    }

    // copy to vec
    v.to_vec()
}

CC @CGMossa i used the &mut [i32] trick i learned from you. Any enhancements you want to make?

Worth noting that you can only call this function once lol because it modifies the vector in place