Question

假设输入如下所示。

indexSet1 = [0,1,2]
indexSet2 = [1,2]
A = [[1,2,3],[4,5,6],[7,8,9]]

然后我想得到一个高度为3且宽度为2的矩阵，元素对应于indexSet1的值和indexSet2的一个。简而言之，我希望通过indexSet1和indexSet2从A中获得数组[[2,3]，[4,5]，[7,8]]。

当我编码如下，我无法得到我想要的结果。

>>> import numpy as np
>>> np.array(A)[np.array(indexSet1),np.array(indexSet2)]
array([5, 9])

任何人都可以知道明智的方法吗？抱歉我的英语不好。并提前感谢你。

Answer 1

使用nested list comprehension：

>>> indexSet1 = [0,1,2]
>>> indexSet2 = [1,2]
>>> A = [[1,2,3],[4,5,6],[7,8,9]]
>>> [[A[i][j] for j in indexSet2] for i in indexSet1]
[[2, 3], [5, 6], [8, 9]]

Answer 2

在NumPy中你可以这样做：

>>> A = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> A[np.array(indexSet1)[:, None], indexSet2]
array([[2, 3],
       [5, 6],
       [8, 9]])

使用索引集从多维数组中提取特定数组

2 个答案: