Использование элементов генератора выборочно

Question

Допустим, у меня есть несколько массивов / списков, которые содержат много значений, что означает, что загрузка нескольких из них в память в конечном итоге приведет к ошибке памяти из-за нехватки памяти. Один из способов обойти это - загрузить эти массивы / списки в генератор, а затем использовать их при необходимости. Однако с генераторами у вас не так много контроля, как с массивами / списками - и это моя проблема.

Позвольте мне объяснить.

В качестве примера у меня есть следующий код, который создает генератор с несколькими небольшими списками. Так что да, это совсем не интенсивно использует память, просто пример:

import numpy as np

np.random.seed(10)

number_of_lists = range(0, 5)

generator_list = (np.random.randint(0, 10, 10) for i in number_of_lists)

Если я перебираю этот список, я получаю следующее:

for i in generator_list:
    print(i)

>> [9 4 0 1 9 0 1 8 9 0]
>> [8 6 4 3 0 4 6 8 1 8]
>> [4 1 3 6 5 3 9 6 9 1]
>> [9 4 2 6 7 8 8 9 2 0]
>> [6 7 8 1 7 1 4 0 8 5]

То, что я хотел бы сделать, это суммировать элементы для всех списков (axis = 0). Таким образом, вышесказанное в свою очередь должно привести к:

[36, 22, 17, 17, 28, 16, 28, 31, 29, 14]

Для этого я мог бы использовать следующее:

sum = [0]*10
for i in generator_list:
    sum += i

где 10 - длина одного из списков.

Пока все хорошо. Я не уверен, есть ли лучший / более оптимизированный способ сделать это, но он работает.

Моя проблема в том, что я хотел бы определить, какие списки в generator_list я хочу использовать. Например, что, если я хотел бы сложить два из первого [0] списка, один из третьего и 2 из последнего, т.е.

[9 4 0 1 9 0 1 8 9 0]
[9 4 0 1 9 0 1 8 9 0]
[4 1 3 6 5 3 9 6 9 1]
[6 7 8 1 7 1 4 0 8 5]
[6 7 8 1 7 1 4 0 8 5]

>> [34, 23, 19, 10, 35, 5, 19, 22, 43, 11]

Как бы я поступил так?

И прежде чем возникнут вопросы, почему я хочу сделать это таким образом, причина в том, что в моем реальном случае получение массивов в генератор занимает некоторое время. Тогда я в принципе мог бы просто сгенерировать новый генератор, в котором я бы упорядочил списки, как показано в новом списке, но опять же, это означало бы, что мне придется подождать, чтобы получить их в новом генераторе. И если это случится тысячи раз (как видно при начальной загрузке), ну, это займет некоторое время. С первым генератором у меня есть ВСЕ списки, которые доступны. Теперь я просто хочу использовать их выборочно, поэтому мне не нужно создавать новый генератор каждый раз, когда я хочу смешать его, и суммировать новый набор массивов / списков.

Answer 1

import numpy as np
np.random.seed(10)

number_of_lists = range(5)

generator_list = (np.random.randint(0, 10, 10) for i in number_of_lists)

indices = [0, 0, 2, 4, 4]
assert sorted(indices) == indices, "only works for sorted list"

# sum_ = [0] * 10

# I prefer this:
sum_ = np.zeros((10,), dtype=int)

generator_index = -1

for index in indices:
    while generator_index < index:
        vector = next(generator_list)
        generator_index += 1
    sum_ += vector

print(sum_)

выходы

[34 23 19 10 37  5 19 22 43 11]