CODEWITHSUNDEEP
pythonarraysnumpyindexingslice
amzon-ex
amzon-ex

Reputation: 1744

Reshape jagged array and fill with zeros

The task I wish to accomplish is the following: Consider a 1-D array a and an array of indices parts of length N. Example:

a = np.arange(9)
parts = np.array([4, 6, 9])

# a = array([0, 1, 2, 3, 4, 5, 6, 7, 8])

I want to cast a into a 2-D array of shape (N, <length of longest partition in parts>), inserting values of a upto each index in indx in each row of the 2-D array, filling the remaining part of the row with zeroes, like so:

array([[0, 1, 2, 3],
       [4, 5, 0, 0],
       [6, 7, 8, 0])

I do not wish to use loops. Can't wrap my head around this, any help is appreciated.

Upvotes: 1

Views: 600

Answers (1)

Divakar
Divakar

Reputation: 221614

Here's one with boolean-indexing -

def jagged_to_regular(a, parts):
    lens = np.ediff1d(parts,to_begin=parts[0])
    mask = lens[:,None]>np.arange(lens.max())
    out = np.zeros(mask.shape, dtype=a.dtype)
    out[mask] = a
    return out

Sample run -

In [46]: a = np.arange(9)
    ...: parts = np.array([4, 6, 9])

In [47]: jagged_to_regular(a, parts)
Out[47]: 
array([[0, 1, 2, 3],
       [4, 5, 0, 0],
       [6, 7, 8, 0]])

Upvotes: 1

Related Questions