從一維列表 numpy 創建一個多維數組

Question

努力用文字描述這個問題，但有一個看似簡單的問題我找不到答案。

我想使用一個列表/數組中的值和另一個列表/數組中的索引創建一個數組。 我希望新數組的形狀與索引數組相同。

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 循環的情況下執行此操作，因為我們正在查看數十萬行並且需要縮短計算時間。 再次感謝 ：）

Answer 1

您可以使用列表理解：

>>> [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 在幕后緩存數據），這可能對您在生產中更有用。 在本地進行測試，看看您是否得到類似或不同的答案。

Answer 2

只需將每個索引子列表應用於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])]

子列表的長度不同，因此無法將結果制成二維數組 - 它必須保留為列表（或者如果您堅持使用一維對象 dtype 數組）。