Apply functions over (anti-)diagonals

user_10068985 0
Viewed 78

We know that apply() in base can apply a function to margins of an array, and margins should be rows or columns. I want to expand the margins to "diagonals" and "anti-diagonals". The structure looks like

diagApply <- function(x, FUN, ..., anti = FALSE) {  }

I name this function diagApply and the argument anti is an indicator of "diagonals" and "anti-diagonals" which FUN will be applied over (defaults to FALSE). Take a 3x4 matrix for example:

mat <- matrix(letters[1:12], 3, 4)

#      [,1] [,2] [,3] [,4]
# [1,] "a"  "d"  "g"  "j" 
# [2,] "b"  "e"  "h"  "k" 
# [3,] "c"  "f"  "i"  "l"

Assume that I need to paste elements in each (anti-)diagonal. With an apply()-like way, the function should be executed as

diagApply(mat, paste, collapse = ".")              # code 1
diagApply(mat, paste, collapse = ".", anti = TRUE) # code 2

The expected output will be

list("c", "b.f", "a.e.i", "d.h.l", "g.k", "j")  # output of code 1
list("a", "d.b", "g.e.c", "j.h.f", "k.i", "l")  # output of code 2

Are there any existing functions achieving that? If not, hope someone can share your idea on this issue. Thanks in advance!

apply r function
2 Answers

My solution is quite similar to the one in @MartinGal comment, only using the function rev and the + operator between row(mat) and col(mat) for the anti-diagonal case.

Such a function should do the job

diagApply <- function(x, FUN, ..., anti = FALSE) {
  
  if(anti == FALSE)
    res <- unname(lapply(split(mat, col(mat) - row(mat)), FUN, ...))
  
  else
    res <- unname(lapply(lapply(split(mat, col(mat) + row(mat)), rev), FUN, ...))

  return(res)
}

The results are indeed equal to your expected output

identical(diagApply(mat, paste, collapse = "."),
          list("c", "b.f", "a.e.i", "d.h.l", "g.k", "j"))
# [1] TRUE
identical(diagApply(mat, paste, collapse = ".", anti = TRUE),
          list("a", "d.b", "g.e.c", "j.h.f", "k.i", "l"))
# [1] TRUE
user_74654620

Building on the other comments, you can do something more general to work from any corner:

diagApply(mat, 1, toString)
# [1] "c"       "b, f"    "a, e, i" "d, h, l" "g, k"    "j"      
diagApply(mat, 2, toString)
# [1] "a"       "b, d"    "c, e, g" "f, h, j" "i, k"    "l"      
diagApply(mat, 3, toString)
# [1] "j"       "g, k"    "d, h, l" "a, e, i" "b, f"    "c"      
diagApply(mat, 4, toString)
# [1] "l"       "i, k"    "f, h, j" "c, e, g" "b, d"    "a"

Your desired results can be obtained by

diagApply(mat, 1, paste, collapse = '.', SIMPLIFY = FALSE)
diagApply(mat, 2, paste, collapse = '.', SIMPLIFY = FALSE)

I'm not sure how important the ordering is for the elements, but you can pass a rev to FUN if desired.

diagApply <- function(X, MARGIN, FUN, ..., SIMPLIFY = TRUE) {
  idx <- switch(
    MARGIN,
    col(X) - row(X),
    row(X) + col(X),
    (col(X) - row(X)) * -1L,
    (row(X) + col(X)) * -1L,
    stop('\'MARGIN\' should be 1, 2, 3, or 4')
  )
  
  res <- unname(lapply(split(X, idx), FUN, ...))
  
  if (!isFALSE(SIMPLIFY) && length(res)) 
    simplify2array(res, higher = (SIMPLIFY == 'array'))
  else res
}
user_29949490
Related