Создайте многомерный массив из 1-D списка, numpy
Пытаюсь описать эту проблему словами, но у меня есть, казалось бы, простая проблема, на которую я не могу найти ответа.
Я хочу создать массив, используя значения из одного списка / массива и индексы из другого. Я хочу, чтобы форма нового массива была такой же, как у индексного массива.
import numpy as np
a = np.array([9, 8, 7, 6, 5, 4, 3, 2, 1, 0, -1, -2])
b = [[0, 1], [2, 3, 4], [6, 7, 8, 9, 10]]
result = func(a, b) #some function or operator...
print(result)
[[9, 8], [7, 6, 5], [3, 2, 1, 0, -1]]
Спасибо! :)
Обновлено:
Пока что хорошие решения, но я бы предпочел сделать это без цикла for, поскольку мы рассматриваем сотни тысяч строк и нам нужно сокращать время вычислений. Еще раз спасибо :)
Ответы 2
Вы можете использовать понимание списка:
>>> [a[x[0]:x[-1]+1] for x in b]
[array([9, 8]), array([7, 6, 5]), array([ 3, 2, 1, 0, -1])]
Обновлено: Ваш вопрос указывает на то, что вам нужен более быстрый вариант, поэтому вы можете протестировать следующий скрипт, чтобы узнать, какой из них быстрее для вашей установки Python:
#!/usr/bin/env python
import timeit
setup = '''
import numpy as np
a = np.array([9, 8, 7, 6, 5, 4, 3, 2, 1, 0, -1, -2])
b = [[0, 1], [2, 3, 4], [6, 7, 8, 9, 10]]
'''
test1 = '''
def test():
return [a[x[0]:x[-1]+1] for x in b]
'''
test2 = '''
def test():
return [a[idx] for idx in b]
'''
print(timeit.timeit(setup = setup,
stmt = test1,
number = 1000000))
print(timeit.timeit(setup = setup,
stmt = test2,
number = 1000000))
На моей машине два представленных вам подхода работают примерно одинаково, но hpaulj отвечать может быть немного быстрее (если Python не кэширует данные за кулисами), что может быть более полезно для вас в продакшене. Протестируйте его на месте и посмотрите, получите ли вы аналогичный или другой ответ.
Просто примените каждый подсписок индексации к a:
In [483]: a = np.array([9, 8, 7, 6, 5, 4, 3, 2, 1, 0, -1, -2])
...:
...: b = [[0, 1], [2, 3, 4], [6, 7, 8, 9, 10]]
...:
...:
In [484]: [a[idx] for idx in b]
Out[484]: [array([9, 8]), array([7, 6, 5]), array([ 3, 2, 1, 0, -1])]
Подсписки различаются по длине, поэтому результат не может быть преобразован в 2-мерный массив - он должен оставаться списком (или, если вы настаиваете на массиве 1d-объекта dtype).