python - 对于一个数组中的每个标签，在另一个数组中将前 k 次出现设置为 False

Question

我有两个(排序的)数组 A 和 B，它们的长度不同，每个数组都包含重复多次的唯一标签。A 中每个标签的计数小于或等于 B 中的计数。A 中的所有标签将出现在 B 中，但 B 中的某些标签不会出现在 A 中。

我需要一个与 B 长度相同的对象，其中对于 A 中的每个标签 i(出现 k_i 次)，第一个 k_i B 中出现的标签 i 需要设置为 False。其余元素应为 True。

下面的代码给出了我所需要的，但是如果 A 和 B 很大，这可能需要很长时间:

import numpy as np

# The labels and their frequency
A = np.array((1,1,2,2,3,4,4,4))
B = np.array((1,1,1,1,1,2,2,3,3,4,4,4,4,4,5,5))

A_uniq, A_count = np.unique(A, return_counts = True)
new_ind = np.ones(B.shape, dtype = bool)
for i in range(len(A_uniq)):
    new_ind[np.where(B == A_uniq[i])[0][:A_count[i]]] = False

print(new_ind)
#[False False  True  True  True False False False  True False False False
#  True  True  True  True]

有没有更快或更有效的方法来做到这一点？我觉得我可能遗漏了一些明显的广播或矢量化解决方案。

Answer 1

这是一个带有 np.searchsorted 的 -

idx = np.searchsorted(B, A_uniq)
id_ar = np.zeros(len(B),dtype=int)
id_ar[idx] = 1
id_ar[A_count+idx] -= 1
out = id_ar.cumsum()==0

我们可以进一步优化计算 A_uniq,A_count 使用它的排序性质而不是使用 np.unique，就像这样 -

mask_A = np.r_[True,A[:-1]!=A[1:],True]
A_uniq, A_count = A[mask_A[:-1]], np.diff(np.flatnonzero(mask_A))