CODEWITHSUNDEEP
rmatlabnumpy
Samuel Isaacson
Samuel Isaacson

Reputation: 365

Indexing a vector by an array in R

In MATLAB and numpy, you can index a vector by an array of indices and get a result of the same shape out, e.g.

A = [1 1 2 3 5 8 13];
B = [1 2; 2 6; 7 1; 4 4];
A(B)

## ans =
##  
##     1    1
##     1    8
##    13    1
##     3    3

or

import numpy as np
a = np.array([1, 1, 2, 3, 5, 8, 13])
b = np.reshape(np.array([0, 1, 1, 5, 6, 0, 3, 3]), (4, 2))
a[b]

## array([[ 1,  1],
##        [ 1,  8],
##        [13,  1],
##        [ 3,  3]])

However, in R, indexing a vector by an array of indices returns a vector:

a <- c(1, 1, 2, 3, 5, 8, 13)
b <- matrix(c(1, 2, 7, 4, 2, 6, 1, 4), nrow = 4)
a[b]

## [1]  1  1 13  3  1  8  1  3

Is there an idiomatic way in R to perform vectorized lookup that preserves array shape?

Upvotes: 2

Views: 142

Answers (3)

KFB
KFB

Reputation: 3501

Option 1: if we do not need to keep the original values in b, we could simply

"Caveat: the values in b will be over-written"
b[] = a[b]
b
#      [,1] [,2]
# [1,]    1    1
# [2,]    1    8
# [3,]   13    1
# [4,]    3    3

Option 2: if want to retain the values in b, An easy workaround could be

c = b  # copy b to c
c[] = a[c]
c
#      [,1] [,2]
# [1,]    1    1
# [2,]    1    8
# [3,]   13    1
# [4,]    3    3

Actually I found Option 2 is easy to follow and clean.

Upvotes: 1

BrodieG
BrodieG

Reputation: 52687

You can't specify dimensions through subsetting alone in R (AFAIK). Here is a workaround:

`dim<-`(a[b], dim(b))

Produces:

     [,1] [,2]
[1,]    1    1
[2,]    1    8
[3,]   13    1
[4,]    3    3

dim<-(...) just allows us to use the dimension setting function dim<- for its result rather than side effect as is normally the case.

You can also do stuff like:

t(apply(b, 1, function(idx) a[idx]))

but that will be slow.

Upvotes: 2

John Paul
John Paul

Reputation: 12684

This is not very elegant, but it works

matrix(a[b],nrow=nrow(b))

Upvotes: 2

Related Questions