使用Python中的單維數組進行多維數組索引

Question

我有二維數組， X大小(500,10)和單維索引數組Y其大小為500其每個條目是相應X行的正確值列的索引，例如， y(0)是2然后它表示X的第一行的第2列是正確的，類似地， y(3) = 4表示X行3和第4列具有正確的值。

我想使用索引數組Y從X獲取所有正確的值，而不使用任何循環，即使用向量化，在這種情況下，輸出應為(500,1) 。 但是當我做X[:,y]然后它給出輸出(500,500) 。 有人可以幫助我如何使用Y，PLZ正確索引數組X.

謝謝大家的幫助。

Answer 1

另一種選擇是多維列表位置索引：

import numpy as np

ncol = 10  # 10 in your case
nrow = 500  # 500 in your case
# just creating some test data:
x = np.arange(ncol*nrow).reshape(nrow,ncol)
y = (ncol * np.random.random_sample((nrow, 1))).astype(int)

print(x)
print(y)
print(x[np.arange(nrow),y.T].T)

這里解釋了語法。 您基本上需要每個維度的索引數組。 在第一個維度中，在您的情況下，這只是[0，...，500]，第二個維度是您的y數組。 我們需要轉置它（.T），因為它必須具有與第一個和輸出數組相同的形狀。 第二個換位不是真的需要，但給你你想要的形狀。

編輯：

性能問題出現了，我嘗試了迄今為止提到的三種方法。 你需要line_profiler來運行以下命令

kernprof -l -v tmp.py

其中tmp.py是：

import numpy as np

@profile
def calc(x,y):
    z = np.arange(nrow)
    a = x[z,y.T].T  # mine, with the suggested speed up
    b = x[:,y].diagonal().T  # Christoph Terasa
    c = np.array([i[j] for i, j in zip(x, y)])  # tobias_k

    return (a,b,c)

ncol = 5  # 10 in your case
nrow = 10  # 500 in your case

x = np.arange(ncol*nrow).reshape(nrow,ncol)
y = (ncol * np.random.random_sample((nrow, 1))).astype(int)

a, b, c = calc(x,y)
print(a==b)
print(b==c)

我的python 2.7.6的輸出：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    3                                           @profile
    4                                           def calc(x,y):
    5         1            4      4.0      0.1      z = np.arange(nrow)
    6         1           35     35.0      0.8      a = x[z,y.T].T
    7         1         3409   3409.0     76.7      b = x[:,y].diagonal().T
    8       501          995      2.0     22.4      c = np.array([i[j] for i, j in zip(x, y)])
    9                                           
    10         1            1      1.0      0.0      return (a,b,c)

其中％Time或Time是相關列。 我不知道如何描述內存消耗，其他人則必須這樣做。 現在看起來我的解決方案對於所請求的尺寸來說是最快的。

Answer 2

雖然從句法角度來看並不是很直觀

X[:,Y].diagonal()[0]

將為您提供您正在尋找的價值。 花式索引從每行選擇X中的所有值Y ，和diagonal只有那些在其中i == j中的索引選擇。 最后用[0]索引只是使2d數組變平。

Answer 3

您需要輔助向量R來索引行

In [50]: X = np.arange(24).reshape((6,4))

In [51]: Y = np.random.randint(0,4,6)

In [52]: R = np.arange(6)

In [53]: Y
Out[53]: array([0, 2, 2, 0, 1, 0])

In [54]: X[R,Y]
Out[54]: array([ 0,  6, 10, 12, 17, 20])

為您的用例

X_y = X[np.arange(500), Y]

---

編輯

我忘了提一下，如果你想要2D結果，你可以使用虛擬索引獲得這樣的結果

X_y_2D = X[np.arange(500), Y, None]