Numpy repeat for 2d array
<p>Given two arrays, say </p>

<pre><code>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)
</code></pre>

<p>np.repeat(arr, rep) returns </p>

<pre><code>array([10, 10, 10, 24, 24, 24, 24], dtype=int32)
</code></pre>

<p>Is there any way to replicate this functionality for a set of 2D arrays?</p>

<p>That is given </p>

<pre><code>arr = array([[10, 24, 24, 24,  1, 21,  1, 21,  0,  0],
            [10, 24, 24,  1, 21,  1, 21, 32,  0,  0]], dtype=int32)
rep = array([[3, 2, 2, 0, 0, 0, 0, 0, 0, 0],
            [2, 2, 2, 0, 0, 0, 0, 0, 0, 0]], dtype=int32)
</code></pre>

<p>is it possible to create a function which vectorizes?</p>

<p>PS: The number of repeats in each row need not be the same. I'm padding each result row to ensure that they are of same size.</p>

<pre><code>def repeat2d(arr, rep):
    # Find the max length of repetitions in all the rows. 
    max_len = rep.sum(axis=-1).max()  
    # Create a common array to hold all results. Since each repeated array will have 
    # different sizes, some of them are padded with zero.
    ret_val = np.empty((arr.shape[0], maxlen))  
    for i in range(arr.shape[0]):
        # Repeated array will not have same num of cols as ret_val.
        temp = np.repeat(arr[i], rep[i])
        ret_val[i,:temp.size] = temp
    return ret_val 
</code></pre>

<p>I do know about np.vectorize and I know that it does not give any performance benefits over the normal version.</p>

<p>Another solution similar to @hpaulj's solution:</p>

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

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

python arrays numpy vectorization