Home > Enterprise >  Sum elements with conditions on indices
Sum elements with conditions on indices

Time:07-15

I have 3 vectors:

x <- c(3, 5, 2)
y <- c(3, 2, 1,  1, 2, 3, 4, 5,  4, 5)
z <- c(2, 4, 8, 1, 5)

x is the number of elements in each group.

y gives indices to extract elements from z. The first three indices belong to group 1 (corresponding to first element in x, 3); the next five indices belong to group 2 (the second element in x, 5), and so on.

z is values from which to extract values using indices in y, and then summed by group.

For each element in x, I want to sum the elements in z indexed by corresponding indices in y. For example:

i = 1; indices = 3, 2, 1; sum = 8 4 2

i = 2; indices = 1, 2, 3, 4, 5; sum = 2 4 8 1 5

i = 3; indices = 4, 5; sum = 1 5

I've thought of creating a new vector with the same length as y and elements are from z with their respective indices but this loop lies within a bigger program and it may get too complicated. I'd really appreciate your advices on how to specify conditions on indices to get the sum.

Thanks!

CodePudding user response:

First index z with y to get a vector of elements you want to sum. Then create a group index from x, and use tapply() to sum in each group:

x <- c(3, 5, 2)
y <- c(3, 2, 1, 1, 2, 3, 4, 5, 4, 5)
z <- c(2, 4, 8, 1, 5)

g <- rep(seq_along(x), x)
tapply(z[y], g, sum)
#>  1  2  3 
#> 14 20  6

CodePudding user response:

With a for loop:

s <- c()
for(i in seq(x)){
  if(i == 1){
    idx <- seq(x[i])
  } else {
    idx <- seq(from = cumsum(x)[i - 1]   1, to = cumsum(x)[i])
  }
  s <- c(s, sum(z[y[idx]]))
}

output

> s
# [1] 14 20  6

CodePudding user response:

Another possible solution, in base R:

sapply(unname(split(y, rep(1:length(x), x))), \(x) sum(z[x]))

#> [1] 14 20  6

  • Related