NumPy 数组按索引部分向下填充行

Question

Let's say I have the following (fictitious) NumPy array:

arr = np.array(
    [[1,   2,  3,  4],
     [5,   6,  7,  8],
     [9,  10, 11, 12],
     [13, 14, 15, 16],
     [17, 18, 19, 20],
     [21, 22, 23, 24],
     [25, 26, 27, 28],
     [29, 30, 31, 32],
     [33, 34, 35, 36],
     [37, 38, 39, 40]
    ]
)

And for row indices idx = [0, 2, 3, 5, 8, 9] I'd like to repeat the values in each row downward until it reaches the next row index:

np.array(
    [[1,   2,  3,  4],
     [1,   2,  3,  4],
     [9,  10, 11, 12],
     [13, 14, 15, 16],
     [13, 14, 15, 16],
     [21, 22, 23, 24],
     [21, 22, 23, 24],
     [21, 22, 23, 24],
     [33, 34, 35, 36],
     [37, 38, 39, 40]
    ]
)

Note that idx will always be sorted and have no repeat values. While I can accomplish this by doing something like:

for  start, stop in zip(idx[:-1], idx[1:]):
    for i in range(start, stop):
        arr[i] = arr[start]

# Handle last index in `idx`
start, stop = idx[-1], arr.shape[0]
for i in range(start, stop):
    arr[i] = arr[start]

Unfortunately, I have many, many arrays like this and this can become slow as the size of the array gets larger (in both the number of rows as well as the number of columns) and the length of idx also increases. The final goal is to plot these as a heatmaps in matplotlib , which I already know how to do. Another approach that I tried was using np.tile :

for  start, stop in zip(idx[:-1], idx[1:]):
    reps = max(0, stop - start)
    arr[start:stop] = np.tile(arr[start], (reps, 1))

# Handle last index in `idx`
start, stop = idx[-1], arr.shape[0]
arr[start:stop] = np.tile(arr[start], (reps, 1))

But I am hoping that there's a way to get rid of the slow for-loop .

Answer 1

Try np.diff to find the repetition for each row, then np.repeat :

# this assumes `idx` is a standard list as in the question
np.repeat(arr[idx], np.diff(idx+[len(arr)]), axis=0)

Output:

array([[ 1,  2,  3,  4],
       [ 1,  2,  3,  4],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16],
       [13, 14, 15, 16],
       [21, 22, 23, 24],
       [21, 22, 23, 24],
       [21, 22, 23, 24],
       [33, 34, 35, 36],
       [37, 38, 39, 40]])