Numpy Repeat For 2d Array

Given two arrays, say arr = array([10, 24, 24, 24, 1, 21, 1, 21, 0, 0], dtype=int32) rep = array([3, 2, 2, 0, 0, 0, 0, 0, 0, 0], dtype=int32) np.repeat(arr, rep) returns arr

Solution 1:

So you have a different repeat array for each row? But the total number of repeats per row is the same?

Just do the repeat on the flattened arrays, and reshape back to the correct number of rows.

In [529]: np.repeat(arr,rep.flat)
Out[529]: array([10, 10, 10, 24, 24, 24, 24, 10, 10, 24, 24, 24, 24,  1])
In [530]: np.repeat(arr,rep.flat).reshape(2,-1)
Out[530]: 
array([[10, 10, 10, 24, 24, 24, 24],
       [10, 10, 24, 24, 24, 24,  1]])

If the repetitions per row vary, we have the problem of padding variable length rows. That's come up in other SO questions. I don't recall all the details, but I think the solution is along this line:

Change rep so the numbers differ:

In [547]: rep
Out[547]: 
array([[3, 2, 2, 0, 0, 0, 0, 0, 0, 0],
       [2, 2, 2, 1, 0, 2, 0, 0, 0, 0]])
In [548]: lens=rep.sum(axis=1)
In [549]: lens
Out[549]: array([7, 9])
In [550]: m=np.max(lens)
In [551]: m
Out[551]: 9

create the target:

In [552]: res = np.zeros((arr.shape[0],m),arr.dtype)

create an indexing array - details need to be worked out:

In [553]: idx=np.r_[0:7,m:m+9]
In [554]: idx
Out[554]: array([ 0,  1,  2,  3,  4,  5,  6,  9, 10, 11, 12, 13, 14, 15, 16, 17])

flat indexed assignment:

In [555]: res.flat[idx]=np.repeat(arr,rep.flat)
In [556]: res
Out[556]: 
array([[10, 10, 10, 24, 24, 24, 24,  0,  0],
       [10, 10, 24, 24, 24, 24,  1,  1,  1]])

Solution 2:

Another solution similar to @hpaulj's solution:

def repeat2dvect(arr, rep):
    lens = rep.sum(axis=-1)
    maxlen = lens.max()
    ret_val = np.zeros((arr.shape[0], maxlen))
    mask = (lens[:,None]>np.arange(maxlen))
    ret_val[mask] = np.repeat(arr.ravel(), rep.ravel())
    return ret_val

Instead of storing indices, I'm creating a bool mask and using the mask to set the values.