Repeating rows in numpy according to a vector of indices

Question

Suppose I have a matrix B:

B = [
    [0, 1, 2],
    [2, 3, 4],
    [5, 6, 7]
]

and a vector a:

a = [0,0,1,1,2]

I need to define a new vector C such that it repeats the rows in B as specified by a, i.e.,

C = [
    [0, 1, 2],
    [0, 1, 2],
    [2, 3, 4]
    [2, 3, 4],
    [5, 6, 7]
]

Is there a trick command to do this in Python?

Answer 1

You can use a as an index array here.

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

And from the docs:

For all cases of index arrays, what is returned is a copy of the original data, not a view as one gets for slices.

In pure Python you can use a list comprehension:

>>> B = [
    [0, 1, 2],
    [2, 3, 4],
    [5, 6, 7]
]
>>> [B[x][:] for x in a]
[[0, 1, 2], [0, 1, 2], [2, 3, 4], [2, 3, 4], [5, 6, 7]]

Note that [:] returns a shallow copy of the lists, if the lists contains mutable objects then you'll have to use copy.deepcopy to get a completely new copy.