Самый быстрый способ повторить список матриц разное количество раз

Question

У меня есть список из N матриц, каждая из которых имеет одинаковую форму (dim1, dim2). У меня есть еще один список из N целых чисел, в котором указано, сколько раз повторять каждую матрицу. Каков самый быстрый способ создать массив numpy каждой матрицы в списке матриц, повторяющийся x раз в соответствии со списком повторений?

Например:

mat_a = np.array([[1, 2],[3, 4]]) # mat_a.shape = (2,2)
mat_b = np.array([[5, 6],[7, 8]]) # mat_b.shape = (2,2)
matrices = [mat_a, mat_b]
repeats = [2, 3]  

result = np.array([mat_a, mat_a, mat_b, mat_b, mat_b]) # results.shape = (5, 2, 2)

Единственный способ, который я мог придумать, - это использовать al oop, что очень медленно:

result = np.array([], dtype=float).reshape(0, 2, 2)
for i in range(len(repeats)):
    result = np.vstack((result, np.tile(matrices[i], (repeats[i], 1, 1))))

Answer 1

Учитывая переменные «матрицы» и «повторения», как вы определили их выше, вам нужно:

result = np.repeat(matrices, repeats, axis=0)

Answer 2

РЕДАКТИРОВАТЬ : теперь посмотрите гораздо лучшее решение Шалома Рохмана ... Я не буду удалять это, на случай, если все еще есть полезный контент, но, очевидно, np.repeat - это способ go.

Это помогает? Если матрицы представляют собой массив numpy:

matrices = np.array([mat_a, mat_b])

И у вас есть список [0, 0, 1, 1, 1]

indices = [i for i, n in enumerate(repeats) for _ in range(n)]

, вы можете проиндексировать его, используя:

matrices[indices]

Что дает:

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

       [[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]],

       [[5, 6],
        [7, 8]],

       [[5, 6],
        [7, 8]]])