CODEWITHSUNDEEP
rcumulative-sum
d.b
d.b

Reputation: 32558

Cumulative sum over matrix diagonals

The input is a square matrix with mostly 0 and some 1. The goal is to take a (sort of) cumulative sum of consecutive 1s along the diagonals of the input matrix.

#Input
ind = rbind(cbind(x = c(2, 3, 1, 2 , 3),
                  y = c(1, 2, 3, 4, 5)))
m1 = replace(matrix(0, 5, 5), ind, 1)
m1
#     [,1] [,2] [,3] [,4] [,5]
#[1,]    0    0    1    0    0
#[2,]    1    0    0    1    0
#[3,]    0    1    0    0    1
#[4,]    0    0    0    0    0
#[5,]    0    0    0    0    0

#Desired Output
#      [,1] [,2] [,3] [,4] [,5]
# [1,]    0    0    0    0    0
# [2,]    0    0    0    0    0
# [3,]    0    2    0    0    3
# [4,]    0    0    0    0    0
# [5,]    0    0    0    0    0

I've got a for loop that does the job but is there a better way?

#Current Approach
m2 = m1
for (i in 2:nrow(m1)){
    for (j in 2:nrow(m1)){
        if (m1[i-1, j-1] == 1 & m1[i, j] == 1){
            m2[i, j] = m2[i - 1, j - 1] + m2[i, j]
            m2[i - 1, j - 1] = 0
        }
    }
}
m2
#     [,1] [,2] [,3] [,4] [,5]
#[1,]    0    0    0    0    0
#[2,]    0    0    0    0    0
#[3,]    0    2    0    0    3
#[4,]    0    0    0    0    0
#[5,]    0    0    0    0    0

Upvotes: 1

Views: 512

Answers (2)

SymbolixAU
SymbolixAU

Reputation: 26258

Your loop in Rcpp

library(Rcpp)

cppFunction('NumericMatrix diagcumsum( NumericMatrix m1 ) {

  int i = 0;
  int j = 0;
  int n_row = m1.nrow();

  NumericMatrix res = Rcpp::clone( m1 );

  for( i = 1; i < n_row; i++ ) {
    for( j = 1; j < n_row; j++ ) {
      if( m1( (i-1), (j-1) ) == 1 && m1( i, j ) == 1 ) {
        res(i, j) = res( (i-1), (j-1) ) + res(i, j);
        res( (i-1), (j-1) ) = 0;
      }
    }
  }
  return res;
}')

diagcumsum( m1 )

#      [,1] [,2] [,3] [,4] [,5]
# [1,]    0    0    0    0    0
# [2,]    0    0    0    0    0
# [3,]    0    2    0    0    3
# [4,]    0    0    0    0    0
# [5,]    0    0    0    0    0

Upvotes: 2

G. Grothendieck
G. Grothendieck

Reputation: 270428

From the example it seems that every diagonal is all zeros or else is a sequence of ones followed by zeros. We assume that that is always the case.

First form a function cum which takes a diagonal x and outputs a vector of zeros the same length except that position sum(x) is to be set to sum(x) .

Then apply that function across diagonals using ave. row(m1)-col(m1) is constant on diagonals and can be used for grouping.

cum <- function(x, s = sum(x)) replace(0 * x, s, s)
ave(m1, row(m1) - col(m1), FUN = cum)

##      [,1] [,2] [,3] [,4] [,5]
## [1,]    0    0    0    0    0
## [2,]    0    0    0    0    0
## [3,]    0    2    0    0    3
## [4,]    0    0    0    0    0
## [5,]    0    0    0    0    0

If the sequence of ones on a diaongal need not start at the beginning of the diagonal but it is still true that there is only one sequence of ones at most on each diagonal then use this in place of cum above:

cum <- function(x, s = sum(x)) replace(0 * x, s + which.max(x) - 1, s)

If there can be more than one sequence of ones on a diagonal then use this in place of cum above:

library(data.table)
cum <- function(x) {
  ave(x, rleid(x), FUN = function(x, s = sum(x)) replace(0 * x, s, s))
}

Upvotes: 6

Related Questions