[英]Using numpy.where() or similar to get specifc values from a row in a matrix
作为MWE,我有一个2d的numpy数组:
import numpy as np
n1 = np.array([[6,7,8,1], [5,2,4,8], [3,4,2,1], [8,7,2,10]])
n1
array([[ 6, 7, 8, 1],
[ 5, 2, 4, 8],
[ 3, 4, 2, 1],
[ 8, 7, 2, 10]])
我想获取第一行中出现“ 6”和“ 8”的索引; 第二行中出现“ 2”和“ 8”; 其中“ 3”和“ 4”出现在第三行中,而“ 7”和“ 2”出现在最后一行中。 也就是说,我有一个numpy数组的列表:
list1 = [np.array([6, 8]), np.array([2, 8]), np.array([3, 4]), np.array([7, 2])]
我想要[n1[i, np.where(ar1)] for i in range(len(list1))]
或类似的东西返回一个新列表,其中包含出现list1的值的列:
returned_list = [np.array([0, 2]), np.array([1, 3]), np.array([0, 1]), np.array([1, 2])]
显然,我已经尝试过[n1[i, np.where(ar1)] for i in range(len(list1))]
。 有任何想法吗?
这是使用NumPy broadcasting
进行数组输出的一种-
In [117]: arr_list1 = np.array(list1)
In [118]: mask = (n1[:,:,None] == arr_list1[:,None,:]).any(2)
In [119]: np.where(mask)[1].reshape(-1,2)
Out[119]:
array([[0, 2],
[1, 3],
[0, 1],
[1, 2]])
说明
基本上,我们通过引入长度为1
的None/np.newaxis
单身arr_list1
/ arr_list1
来将n1
和arr_list1
到3D
数组,这样当彼此比较时,将与这两个数组中的最后一个轴进行元素比较3D
阵列。
然后,我们沿着最后一个轴查找ANY
匹配项,该匹配项的长度为2
对应于每一行的arr_list1
的两个arr_list1
。 这给了我们2D阵列。 最后,我们需要匹配的行索引,所以np.where()[1]
。
分步进行仔细观察-
1)输入:
In [124]: n1
Out[124]:
array([[ 6, 7, 8, 1],
[ 5, 2, 4, 8],
[ 3, 4, 2, 1],
[ 8, 7, 2, 10]])
In [125]: arr_list1
Out[125]:
array([[6, 8],
[2, 8],
[3, 4],
[7, 2]])
2)比较:
In [126]: (n1[:,:,None] == arr_list1[:,None,:])
Out[126]:
array([[[ True, False],
[False, False],
[False, True],
[False, False]],
[[False, False],
[ True, False],
[False, False],
[False, True]],
[[ True, False],
[False, True],
[False, False],
[False, False]],
[[False, False],
[ True, False],
[False, True],
[False, False]]], dtype=bool)
3) ANY
减少:
In [127]: (n1[:,:,None] == arr_list1[:,None,:]).any(2)
Out[127]:
array([[ True, False, True, False],
[False, True, False, True],
[ True, True, False, False],
[False, True, True, False]], dtype=bool)
In [128]: mask = _
4)最后得到匹配的行索引:
In [130]: np.where(mask)[1].reshape(-1,2)
Out[130]:
array([[0, 2],
[1, 3],
[0, 1],
[1, 2]])
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.