这是向量化的解决方案:
import numpy as np
startRavel = np.array([1,2,2,2,3,3])
unique_policies = np.array([1,2,3])
使用np.argsort对startRavel进行排序.
ix = np.argsort(startRavel)
s_startRavel = startRavel[ix]
使用np.searchsorted查找应该在startRavel中插入unique_policies以保持顺序的索引:
s_ix = np.searchsorted(s_startRavel, unique_policies)
# array([0, 1, 4])
然后使用np.split使用获得的索引拆分数组. s_ix再次使用np.argsort来处理未排序的输入:
ix_r = np.argsort(s_ix)
ixs = np.split(ix, s_ix[ix_r][1:])
np.array(ixs)[ix_r]
# [array([0]), array([1, 2, 3]), array([4, 5])]
?通用解决方案:
让我们将其包装在一个函数中:
def ix_intersection(x, y):
"""
Finds the indices where each unique
value in x is found in y.
Both x and y must be numpy arrays.
----------
x: np.array
Must contain unique values.
Values in x are assumed to be in y.
y: np.array
Returns
-------
Array of arrays. Each array contains the indices where a
value in x is found in y
"""
ix_y = np.argsort(y)
s = np.searchsorted(y[ix_y], x)
ix_r = np.argsort(s)
ixs = np.split(ix_y, s[ix_r][1:])
return np.array(ixs)[ix_r]
?其他例子
让我们尝试以下数组:
startRavel = np.array([1,3,3,2,2,2])
unique_policies = np.array([1,2,3])
ix_intersection(unique_policies, startRavel)
# array([array([0]), array([3, 4, 5]), array([1, 2])])
另一个例子,这次输入未排序:
startRavel = np.array([1,3,3,2,2,2,5])
unique_policies = np.array([1,2,5,3])
ix_intersection(unique_policies, startRavel)
# array([array([0]), array([3, 4, 5]), array([6]), array([1, 2])])