Extract values from matrix using col row indices

3
votes

Say I have two matrices:

> a
     [,1] [,2] [,3] [,4] [,5] [,6] [,7]
[1,]    6   10    5    7    2    2    6
[2,]   10    6    7    7    4    3   12
[3,]   11   10    2   10    6   11    9

and

> b
         [,1] [,2] [,3]
    [1,]    4    1    4
    [2,]    3    6    3
    [3,]    2    5    2

The number of rows in a and b is identical. I am looking for a vectorized way to extract items from a indicated by the column numbers in b, on a row-by-row basis. The result c should therefore look as follows:

> c
     [,1] [,2] [,3]
[1,]    7    6    7
[2,]    7    3    7
[3,]    10    6   10

a[,b[1,]] or a[,b[2,]] or a[,b[3,]] manage to get the correct results for rows 1, 2 and 3 respectively only. Can this be done with a simple matrix function at all? Is apply necessary?

I have tried to adapt a solution to a similar problem in Index values from a matrix using row, col indicies but didn't understand how cbind was used here to extract matrix elements.

2
rmatrixindexing

2 Answers

3
votes

Here is a cbind version

 t(`dim<-`(a[cbind(rep(1:nrow(a), each=ncol(b)), c(t(b)))], dim(b)))
 #     [,1] [,2] [,3]
 #[1,]    7    6    7
 #[2,]    7    3    7
 #[3,]   10    6   10

Or as suggested by @thelatemail

 matrix(a[cbind(c(row(b)),c(b))],nrow=nrow(a))
 #     [,1] [,2] [,3]
 #[1,]    7    6    7
 #[2,]    7    3    7
 #[3,]   10    6   10

Benchmarks

set.seed(24)
a1 <- matrix(sample(1:10, 2e5*7, replace=TRUE), ncol=7)
set.seed(28)
b1 <- matrix(sample(1:7,2e5*3, replace=TRUE), ncol=3)

f1 <- function() {s <- seq_len(nrow(a1))
 t(vapply(s, function(i) a1[i, b1[i,]],numeric(ncol(b1))))
}
f2 <- function() {matrix(a1[cbind(c(row(b1)),c(b1))], nrow=nrow(a1)) }
f3 <- function(){t(`dim<-`(a1[cbind(rep(1:nrow(a1),
                    each=ncol(b1)), c(t(b1)))], dim(b1)))} 
library(microbenchmark)
microbenchmark(f1(), f2(), f3(), unit='relative', times=10L)
#Unit: relative
# expr       min        lq      mean    median        uq       max neval cld
#f1() 16.636045 16.603856 15.319595 15.799335 13.869147 14.629315    10   b
#f2()  1.000000  1.000000  1.000000  1.000000  1.000000  1.000000    10  a 
#f3()  1.310433  1.306228  1.258715  1.278504  1.237299  1.236448    10  a
3
votes

You could try

t(sapply(seq_len(nrow(a)), function(i) a[i, b[i, ]]))
#      [,1] [,2] [,3]
# [1,]    7    6    7
# [2,]    7    3    7
# [3,]   10    6   10

And you may see a slight speed improvement from the sapply solution above with vapply

s <- seq_len(nrow(a))
t(vapply(s, function(i) a[i, b[i, ]], numeric(ncol(b))))
#      [,1] [,2] [,3]
# [1,]    7    6    7
# [2,]    7    3    7
# [3,]   10    6   10

Or a for loop solution is

m <- matrix(, nrow(b), ncol(b))
for(i in seq_len(nrow(a))) { m[i, ] <- a[i, b[i, ]] }
m
#      [,1] [,2] [,3]
# [1,]    7    6    7
# [2,]    7    3    7
# [3,]   10    6   10