CODEWITHSUNDEEP
pythonarraysnumpy
Bravo
Bravo

Reputation: 677

Repeating rows in numpy according to a vector of indices

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?

Upvotes: 4

Views: 1306

Answers (2)

Carsten
Carsten

Reputation: 18446

Use map() to iterate though all elements of a and get the corresponding lists from B:

C = map(lambda x: B[x], a)

Please note that the elements of C will only be references to the rows of B. This might be the behaviour you want. If you need a real copy of the elements, you could return B[x][:] in the lambda for one more level of copying, or use copy.deepcopy() for a complete copy.

Upvotes: 0

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250891

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.

Upvotes: 5

Related Questions