Python中使用单维数组索引多维数组

Question

Python中使用单维数组索引多维数组

3

我有一个二维数组X，大小为(500,10)，还有一个单维索引数组Y，大小为500，其中每个条目都是相应行的正确值列的索引，例如y(0)是2，则表示第一行的第2列的值是正确的，同样地，y(3) = 4 表示X的第3行和第4列具有正确的值。

我想使用索引数组Y从X中获取所有正确的值，而不使用任何循环，即使用向量化，此时输出应为(500,1)。但当我执行X[:,y]时，它给出了(500,500)的输出。请问有人能帮我正确地使用Y对X进行索引吗？

谢谢大家的帮助。

- samquest

请您添加一些代码并展示您尝试过的内容。 - Trev Davies

3

你正在使用 numpy 吗？像这样的东西基本上就是它的用途。 - Jan Christoph Terasa

我不确定我完全理解您的输入和输出。您能否给出X和Y可能的具体（但截短的）示例，并说明您期望的输出是什么？ - ymbirtt

你为什么想要避免循环呢？看起来这个数据结构需要至少一次对数据的遍历，将正确的值提取到一个新的列表中。 - Aaron D

你是否清楚地了解在Python中数组是从0开始索引的呢？因此，y(3)很可能意味着第4行。 - martineau

3个回答

4

虽然从语法角度来看并不直观，但是需要注意的是：

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

将会为您提供您正在寻找的值。精细索引从X的每一行中选择Y中的所有值，而diagonal仅选择在i == j处的索引处的值。在末尾使用[0]进行索引只是将2d数组展平。

- Jan Christoph Terasa

这确实有效，但我想知道是否比使用循环更快/更高效，因为这仍会暂时创建一个500x500的数组。 - tobias_k

我尝试了一下，对于 500 行数据，一个简单的 np.array([x[y] for x, y in zip(X, Y)]) 似乎要快 5-10 倍。 - tobias_k

@tobias_k 我也进行了性能分析，对我来说zip方法似乎更慢。 - StefanS

这个解决方案不好，因为通过 X[:,Y] 进行索引会创建一个新的数组，这是一个缓慢的过程。@StefanS 的解决方案只会创建一个数组的视图，这要快得多。 - Jan Christoph Terasa

@ChristophTerasa 非常感谢您的帮助。 - samquest

显示剩余2条评论

4

你需要一个辅助向量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]

编辑

我忘了提到，如果你想获得一个二维结果，你可以使用一个虚拟索引来获得这样的结果。

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

- gboffi

这似乎基本上与@StefanS的解决方案相同，只是没有使用两个.T。 - tobias_k

这里不需要第一个“.T”，因为你的Y与我的y形状不同，输出的形状为（1,500），而不是所要求的（500,1）。 - StefanS

@tobias_k，看起来这个问题有一个独特、优雅、明显的解决方案...我花了更多时间（全靠自己）才想到它，但我的实现略有不同，而且解释更简单。 - gboffi

@gboffi 谢谢您的热心帮助。 - samquest

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- StefanS · Accepted Answer

另一种选择是多维位置列表索引：

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是相关列。我不知道如何分析内存消耗，需要别人来完成。目前看来，我的解决方案在所请求的维度中是最快的。