Выберите случайные ненулевые элементы из каждой строки двумерного массива

Question

У меня есть 2d массив

a = array([[5, 0, 1, 0],
           [0, 1, 3, 5],
           [2, 3, 0, 0],
           [4, 0, 2, 4],
           [3, 2, 0, 3]])

и массив 1d

b = array([1, 2, 1, 2, 2])

which (b) сообщает, сколько ненулевых элементов мы хотим выбрать из каждой строки массива a.

Например, b[0] = 1 говорит нам, что мы должны выбрать 1 ненулевой элемент из a[0], b[1] = 2 говорит нам, что мы должны выбрать 2 ненулевых элемента из a[1], и так далее.

Для 1d массива это можно сделать с помощью np.random.choice, но я не могу найти, как это сделать для 2d массива, поэтому я должен использовать цикл for, который замедляет вычисления.

Я хочу получить результат в виде 2d-массива как

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

Здесь у нас есть 1 элемент в строке 1, 2 элемента в строке 2, 1 элемент в строке 3 и т. Д., Как указано в массиве b.

Answer 1

Это похоже на проблему конкурентного программирования.

Я не думаю, что вы можете достичь результатов, используя numpy.random.choice (я могу ошибаться).

В любом случае, думай об этом так. Чтобы выбрать x количество ненулевых элементов из одномерного массива размером n, в худшем случае он будет иметь сложность O (n). И для двумерного массива это будет O (n ^ 2), если вы будете следовать тому же наивному подходу.

этот пост почти аналогичен вашему вопросу, но numpy.nonzero также является функцией O (n ^ 2).