使用 numpy 中的数组清理数组索引答案

【问题标题】：Clean array indexing with arrays in numpy使用 numpy 中的数组清理数组索引
【发布时间】：2018-05-19 02:26:13
【问题描述】：

我一遍又一遍地遇到这个问题，似乎无法找到一个干净的解决方案。所以我试图用另一个数组索引一个数组。我有一个二维 numpy 数组。还有一个 1d numpy 数组，其长度与我尝试索引的 2d 数组的第一个维度相同，元素代表我尝试提取的列的索引：

import numpy as np

A = np.random.rand((5,3))

B = np.asarray([2,1,2,0,1])

我想要的行为是为所有行提取数组 B 中的相应列。这可以通过

C = A[np.arange(A.shape[0]),B]

但我可以想象有更好的方法来获得这种行为。使用 : 作为索引第一行会给出错误的行为。

如果有更清洁的方式来做这件事，那就太好了。我真的习惯了 Matlab 中的这个数组索引，但也许在 numpy 中没有等价物。使用布尔索引当然是一种选择，但这也需要一直转换数组。

【问题讨论】：

将A[np.arange(A.shape[0]),B] 包装在一个函数中并调用它really_clean_solution()?
您将如何在 MATLAB 中编写此代码？ MATLAB 对成对和块索引的处理是不同的。
@Divakar 的非讽刺变体：def rl_(I): return np.arange(I.size) rl 代表 range_like。 A[rl_(B), B] 还不错吧？
说真的，在 MATLAB 中执行此操作的更简洁的方法是什么？我一定是用那种语言过时了，但我找不到它。

标签： python numpy

【解决方案1】：

In [60]: A = np.arange(1,16).reshape(5,3)
In [61]: B = np.array([2,1,2,0,1])
In [62]: C = A[np.arange(A.shape[0]),B]
In [63]: C
Out[63]: array([ 3,  5,  9, 10, 14])

八度音

>>  A = reshape(1:15, 3,5).';
>> B = [3,2,3,1,2];
>> A
A =

    1    2    3
    4    5    6
    7    8    9
   10   11   12
   13   14   15

>> A(:,B)
ans =

    3    2    3    1    2
    6    5    6    4    5
    9    8    9    7    8
   12   11   12   10   11
   15   14   15   13   14

这和numpy一样：

In [65]: A[:,B]
Out[65]: 
array([[ 3,  2,  3,  1,  2],
       [ 6,  5,  6,  4,  5],
       [ 9,  8,  9,  7,  8],
       [12, 11, 12, 10, 11],
       [15, 14, 15, 13, 14]])

您暗示从 MATLAB/Octave 中的每一列索引一个项目是干净的，但我错过了它。我过去经常使用该语言工作，但我已经脱离实践了。

sub2ind 完成这项工作：

>> sub2ind([3,5],B, 1:5)
ans =

    3    5    9   10   14

>> A.'(:)(sub2ind([3,5],B,1:5))
ans =

    3
    5
    9
   10
   14

（“F”与“C”的排序使我的比较变得复杂）

numpy 有一个类似的ravel_multi_index：

In [69]: np.ravel_multi_index((np.arange(5),B),(5,3))
Out[69]: array([ 2,  4,  8,  9, 13], dtype=int32)
In [71]: A.flat[_]
Out[71]: array([ 3,  5,  9, 10, 14])

【讨论】：

【解决方案2】：

我想你关注的是np.choose(B,A.T)：

In [125]: A
Out[125]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

In [126]: B = np.asarray([2,1,2,0,1])

In [127]: np.choose(B,A.T)
Out[127]: array([ 2,  4,  8,  9, 13])

【讨论】：

我不建议这样做。试试A = np.zeros((2, 32))、B = np.zeros((2,), int)。
choose 的最大尺寸为 32
np.choose 文档不鼓励这种使用，请参阅注释，“choices 既不应该，也不应该被认为是单个数组，即，最外层的类似序列的容器应该是列表或元组。”