CODEWITHSUNDEEP
pythonnumpy
Ben S.
Ben S.

Reputation: 3545

Make a new array selectively copying, with varying multiplicity, rows of old array in numpy

I have a 2D array with n rows in numpy, and an accompanying 1-D array of n elements in which the ith element specifies how many times the ith row of the original array should go into a new array. For example if my 2D array is:

array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5]])

And my 1D array is 

array([2, 0, 1, 0, 3])

Then I'd like the new array to be:

array([[1, 1, 1],
       [1, 1, 1],
       [3, 3, 3],
       [5, 5, 5],
       [5, 5, 5],
       [5, 5, 5]])

I can't figure out how to do this efficiently, does anyone have any ideas?

Upvotes: 2

Views: 276

Answers (2)

Divakar
Divakar

Reputation: 221514

One approach based on np.cumsum -

import numpy as np

A = np.array([[1, 1, 1],          # Sample Input Array
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5]])

lens = np.array([2, 0, 1, 0, 3])  # Sample Lengths of replications

mask = lens!=0                    # Mask of valid lengths
clens = np.cumsum(lens[mask])     # CUMSUM masked lenghts

# Setup Row-ID array for replications
id = np.zeros((1,clens[-1]),int).ravel()
id[clens[:-1]] = 1

# Finally perform replications for the final output
out = A[mask][id.cumsum(),:]

Upvotes: 1

Zero
Zero

Reputation: 76917

You could use np.repeat() to repeat elements of an array.

In [174]: x.repeat(np.array([2, 0, 1, 0, 3]), axis=0)
Out[174]:
array([[1, 1, 1],
       [1, 1, 1],
       [3, 3, 3],
       [5, 5, 5],
       [5, 5, 5],
       [5, 5, 5]])

Details:

In [175]: x
Out[175]:
array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5]])

In [176]: repeat_on = np.array([2, 0, 1, 0, 3])

In [177]: x.repeat(repeat_on, axis=0)
Out[177]:
array([[1, 1, 1],
       [1, 1, 1],
       [3, 3, 3],
       [5, 5, 5],
       [5, 5, 5],
       [5, 5, 5]])

Upvotes: 8

Related Questions