Is there a function in R that splits a matrix along a margin using a factor or categorical vector?

Viewed 308

I am trying to split a matrix in the same way that you might split a data.frame using split. Is there a function that does that? For example, I have matrix m and I am trying to split it into a list of matrices using vector g.

m <- matrix(rnorm(50), ncol = 5)
groups <- c('A', 'B', 'C')
g <- sample(groups, 10, replace = T)

split doesn't seem to work with matrices so we could convert it into a data.frame:

split(data.frame(m), f = g)

This works but I'd like to keep it as a matrix. The following loop works:

lapply(groups, function(x) m[g == x,])

But is there a dedicated function, or a better way?

4 Answers

We can split on the sequence of rows of 'm', and use that index to subset the rows of 'm'

lapply(split(seq_len(nrow(m)), g), function(i) m[i,])

Here is a way to split a matrix using lapply/split.

lapply(split(m, g), matrix, ncol = ncol(m))

This can easily be written as a one-line function but I prefer a version with some error check.

mat_split <- function(x, f) {
  stopifnot(nrow(x) == length(f))
  lapply(split(x, f), matrix, ncol = ncol(x))
}

Edit

The original question is (my emphasis):

But is there a dedicated function, or a better way?

Following this comment by user20650 there is a function or, better said, a method.

The split.data.frame method can solve the problem.

split.data.frame(m, g)

And this is written in the documentation. From help('split') (my emphasis).

split and split<- are generic functions with default and data.frame methods. The data frame method can also be used to split a matrix into a list of matrices, and the replacement form likewise, provided they are invoked explicitly.

You could iterate over unique values of g and use it to subset the matrices.

sapply(unique(g), function(x) m[g == x, ])

#$B
#        [,1]   [,2]   [,3]   [,4]   [,5]
#[1,] -0.3752 -0.180 -0.333 -0.634  2.011
#[2,] -0.0628 -0.537  2.089  2.069  0.605
#[3,] -0.7635  1.028 -0.779  1.205 -0.110

#$A
#       [,1]   [,2]  [,3]   [,4]   [,5]
#[1,] -0.128 -1.037 0.512 -1.104  0.303
#[2,] -0.246 -0.691 1.303 -0.571  2.023
#[3,] -1.374  0.867 0.810 -0.904 -0.268
#[4,]  0.111 -0.013 0.827  1.294  0.999

#$C
#      [,1]   [,2]   [,3]   [,4]   [,5]
#[1,] 0.712  1.465 -0.471 -0.383 -0.191
#[2,] 1.400  0.121  0.360 -0.890  0.412
#[3,] 0.967 -1.176  0.146  0.570 -0.143

A couple of solutions have been proposed. Here's a benchmark for a matrix of size 100,000 * 5 being split into 1,000 groups:

nr <- 1e5
m <- matrix(rnorm(nr * 5), ncol = 5)
groups <- seq_len(1000)
g <- sample(groups, nr, replace = T)

microbenchmark(
  data.frame = split(data.frame(m), f = g),
  split.data.frame = split.data.frame(m, g),
  matrix = lapply(split(m, g), matrix, ncol = ncol(m)),
  lapply1 = lapply(groups, function(x) m[g == x,]),
  lappyl2 = lapply(split(seq_len(nrow(m)), g), function(i) m[i,])
)



Unit: milliseconds
             expr       min        lq      mean    median        uq       max neval
       data.frame 101.50167 119.37017 132.39754 124.60196 133.04204  299.9586   100
 split.data.frame  14.82502  17.43736  24.66659  18.96938  25.33538  119.5009   100
           matrix  18.99796  22.73603  28.14735  25.82694  31.52766  102.1667   100
          lapply1 699.65742 778.61135 817.87159 811.95775 840.05130 1089.5721   100
          lappyl2  15.37083  17.58404  24.13295  19.08363  24.65315  106.8594   100

For a small number of groups, lapply1 appears to be the faster method.

Related