将对象数组作为索引传递给二维数组

Question

给定一个二维数组：

a = np.array([[10,0,30,10],[40,50,60,10],[70,80,90,10]])

作为对象数组的索引数组：

i = np.array([[0,1],[0,2],[0,1,2]])  #Note different lengths

预期结果：

e = [[10,0,30,10,40,50,60,10],[10,0,30,10,70,80,90,10],[10,0,30,10,40,50,60,10,70,80,90,10]]   

什么有效：

e = [np.hstack(a[i[j]]) for j in range(len(i))]

有没有办法以纯矢量化的方式做到这一点？

我发现 numpy.where() 不起作用，因为索引数组（即 i）中的元素应该具有相同的长度，而我的情况不是。有人能指出我正确的方向吗？

编辑： 除了上述问题，我还想知道当数组“a”更改为时如何执行相同的操作：

a = np.array([[10,0,30,10],[40,50,60,10],[70,80,90,10,30]])#NOTE:Jagged array

索引数组“i”保持不变！

Answer 1

如果a是锯齿状的，而i是多维的，我们可以使用i来索引a：

In [78]: a = np.array([[10,0,30,10],[40,50,60,10],[70,80,90,10,30]])#NOTE:Jagged array                 
In [79]: i = np.array([[0,1],[0,2],[1,2]])                                                             

In [80]: a.shape    #  an array of list objects                                                                                       
Out[80]: (3,)

In [81]: a[i]                                                                                          
Out[81]: 
array([[list([10, 0, 30, 10]), list([40, 50, 60, 10])],
       [list([10, 0, 30, 10]), list([70, 80, 90, 10, 30])],
       [list([40, 50, 60, 10]), list([70, 80, 90, 10, 30])]], dtype=object)

由于这些是列表对象，我们可以使用sum 来“连接”它们：

In [82]: a[i].sum(axis=1)                                                                              
Out[82]: 
array([list([10, 0, 30, 10, 40, 50, 60, 10]),
       list([10, 0, 30, 10, 70, 80, 90, 10, 30]),
       list([40, 50, 60, 10, 70, 80, 90, 10, 30])], dtype=object)

你的列表理解：

In [83]: e = [np.hstack(a[i[j]]) for j in range(len(i))]                                               
In [84]: e                                                                                             
Out[84]: 
[array([10,  0, 30, 10, 40, 50, 60, 10]),
 array([10,  0, 30, 10, 70, 80, 90, 10, 30]),
 array([40, 50, 60, 10, 70, 80, 90, 10, 30])]

一些时间安排：

In [85]: timeit a[i].sum(axis=1)                                                                       
8.64 µs ± 17.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [86]: timeit e = [np.hstack(a[i[j]]) for j in range(len(i))]                                        
63.3 µs ± 168 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

您的hstack 可能会更慢，因为它会将a 列表转换为数组。让我们绕过那个：

In [89]: [sum(a[i[j]],[]) for j in range(len(i))]                                                      
Out[89]: 
[[10, 0, 30, 10, 40, 50, 60, 10],
 [10, 0, 30, 10, 70, 80, 90, 10, 30],
 [40, 50, 60, 10, 70, 80, 90, 10, 30]]
In [90]: timeit [sum(a[i[j]],[]) for j in range(len(i))]                                               
8.41 µs ± 109 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

有时纯列表解决方案更快。转换数组列表需要时间。

===

如果两个数组都是均衡且多维的，我们可以使用纯“矢量化”解决方案：

In [104]: aa = np.array([[10,0,30,10],[40,50,60,10],[70,80,90,10]])                                    
In [105]: i                                                                                            
Out[105]: 
array([[0, 1],
       [0, 2],
       [1, 2]])
In [106]: aa[i]                                                                                        
Out[106]: 
array([[[10,  0, 30, 10],
        [40, 50, 60, 10]],

       [[10,  0, 30, 10],
        [70, 80, 90, 10]],

       [[40, 50, 60, 10],
        [70, 80, 90, 10]]])
In [107]: aa[i].reshape(3,-1)                                                                          
Out[107]: 
array([[10,  0, 30, 10, 40, 50, 60, 10],
       [10,  0, 30, 10, 70, 80, 90, 10],
       [40, 50, 60, 10, 70, 80, 90, 10]])
In [108]: timeit aa[i].reshape(3,-1)                                                                   
5.07 µs ± 57.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

但是，一旦一个或多个数组/列表出现问题，您就会失去此选项，需要认真考虑列表替代方案。

Answer 2

Numpy 数组倾向于具有相同的大小。在您的情况下，您的 i 数组将与列表列表一样工作，因为长度不一致。

你的答案很好，我有另一个解决方案，但它的工作方式与你的类似：

import numpy as np

a = np.array([[10,0,30,10],[40,50,60,10],[70,80,90,10]])
# 'i' can also be a list of lists
i = [[0,1],[0,2],[0,1,2]]

[np.concatenate(a[j]) for j in i]
>>> [array([10,  0, 30, 10, 40, 50, 60, 10]),
     array([10,  0, 30, 10, 70, 80, 90, 10]),
     array([10,  0, 30, 10, 40, 50, 60, 10, 70, 80, 90, 10])]

另一种使用向量的方法：

# make 'i' a consistent shape by including NaN values 
i = np.array([[0,1,np.nan],[0,2,np.nan],[0,1,2]])

# filter out NaN values and index
a[i[np.isfinite(i)].astype(int)]
>>> array([[10,  0, 30, 10],
           [40, 50, 60, 10],
           [10,  0, 30, 10],
           [70, 80, 90, 10],
           [10,  0, 30, 10],
           [40, 50, 60, 10],
           [70, 80, 90, 10]])

我认为此时您仍然必须使用循环来正确连接，因此第一种方法或您的答案更直接。

Answer 3

如果您只是想隐藏for 循环，您可以像这样使用map：

In [891]: import numpy as np

In [892]: a = np.array([[10,  0, 30, 10], 
     ...:               [40, 50, 60, 10], 
     ...:               [70, 80, 90, 10]])

In [893]: i = [[0, 1], 
     ...:      [0, 2], 
     ...:      [0, 1, 2]]

In [894]: e = [[10, 0, 30, 10, 40, 50, 60, 10], 
     ...:      [10, 0, 30, 10, 70, 80, 90, 10], 
     ...:      [10, 0, 30, 10, 40, 50, 60, 10, 70, 80, 90, 10]]

In [895]: e == list(map(lambda x: a[x].flatten().tolist(), i))
Out[895]: True

请注意，上面的代码实际上并未矢量化（请参阅List comprehension vs. map）。