Question

给定一组数据，有没有办法通过N维索引列表查询数据？

示例：

import numpy as np
data = np.array([[-14., 2.,  19.],
                 [-13., 1.,  20.],
                 [-15., 2.,  18.],
                 [-13., 0.,  19.],
                 [-15., 1.,  19.],
                 [-14., 0.,  19.],
                 [-14., 1.,  20.]])


# Uniformly shaped array: works
queries = np.array([[2, 4, 6, 0], [3, 6, 4, 5]])
print data[queries]

# Properly returns
#[[[-15.   2.  18.]
#  [-15.   1.  19.]
#  [-14.   1.  20.]
#  [-14.   2.  19.]]
#
# [[-13.   0.  19.]
#  [-14.   1.  20.]
#  [-15.   1.  19.]
#  [-14.   0.  19.]]]


# N-dimentional array fails
queries = np.array([[4, 6, 0], [3, 6, 4, 5]])
print data[queries]

# IndexError: arrays used as indices must be of integer (or boolean) type #
#
# Desired result:
#[[[-15.   1.  19.]
#  [-14.   1.  20.]
#  [-14.   2.  19.]]
#
# [[-13.   0.  19.]
#  [-14.   1.  20.]
#  [-15.   1.  19.]
#  [-14.   0.  19.]]]

Answer 1

查询中的两个元素具有不同的长度，因此它们存储为列表而不是numpy数组;类似地，结果也将在内部存储为列表，并且不再使用numpy数组对python列表的优势;你可以做的最好的事情是正常的循环：

[data[query].tolist() for query in queries]

#[[[-15.0, 1.0, 19.0], 
#  [-14.0, 1.0, 20.0], 
#  [-14.0, 2.0, 19.0]],
#
# [[-13.0, 0.0, 19.0],
#  [-14.0, 1.0, 20.0],
#  [-15.0, 1.0, 19.0],
#  [-14.0, 0.0, 19.0]]]

或者如果你想将结果部分保留为numpy数组：

[data[query] for query in queries]

#[array([[-15.,   1.,  19.],
#        [-14.,   1.,  20.],
#        [-14.,   2.,  19.]]), array([[-13.,   0.,  19.],
#        [-14.,   1.,  20.],
#        [-15.,   1.,  19.],
#        [-14.,   0.,  19.]])]

Answer 2

print [data[q] for q in queries]

使用N维索引列表查询numpy数组

2 个答案: