如何从数组B中筛选出数组A各连续序列的首个元素？

现有一个分段升序的有序整数数组A：

A = np.array([1,2,3,4, 7,8,9, 23,24,25, 100])

另有数组B，元素从A中选取：

B = np.array([1,2,23,25,100])

需求是过滤B中属于A同一连续序列的元素，仅保留每个唯一序列的首个元素，预期结果：

C = np.array([1,23,100])

目前通过维护追踪列表实现，但方法繁琐，询问是否有更优解法。现有实现代码如下：

import numpy as np

A = np.array([1,2,3,4, 7,8,9, 23,24,25, 100])
B = np.array([1,2,23,25,100])
C = []

already_used_sequence = []

for x in enumerate(A):
    if x[0]-x[1] in already_used_sequence : #did we already group this sequence?
        pass
    elif len(np.intersect1d(B, x[1])) is not None: #is this value in B?
        for h in B:
            if h == x[1]: 
                C.append(x[1])
                already_used_sequence.append(x[0]-x[1])
C=np.array(C)

利用numpy向量运算结合字典映射可以实现更简洁高效的方案，避免原代码中的多层循环和冗余操作：

import numpy as np

A = np.array([1,2,3,4, 7,8,9, 23,24,25, 100])
B = np.array([1,2,23,25,100])

# 为A中每个元素标记所属连续序列的组ID
diff = np.diff(A)
# 从第一个元素开始，每遇到非连续分界（差值>1）就递增组ID
group_ids = np.cumsum(np.concatenate([[0], diff > 1]))

# 建立元素到组ID的映射，方便快速查询
elem_group_map = dict(zip(A, group_ids))

# 遍历B，仅保留每个组的首个出现元素
seen_groups = set()
C = []
for num in B:
    gid = elem_group_map[num]
    if gid not in seen_groups:
        seen_groups.add(gid)
        C.append(num)

C = np.array(C)
print(C)  # 输出: array([  1,  23, 100])

方案说明 #

1. 组ID生成：通过np.diff快速定位A中连续序列的分界点，再用np.cumsum生成唯一组ID，同一连续序列内的元素会被赋予相同ID，时间复杂度O(n)。
2. 快速映射：字典映射让B中元素的组ID查询操作达到O(1)，比原代码中循环查找高效得多。
3. 组去重：用集合记录已处理的组，遍历B时仅保留每个组的第一个元素，逻辑直观且无冗余计算。

相比原实现，该方案代码更简洁，执行效率更高，尤其适合处理大规模数组。

内容的提问来源于stack exchange，提问作者Kelly