在NumPy数组中查找重复序列的索引

时间:2020-01-09 15:31:52

标签: python numpy

这是对previous question的跟进。如果我有一个NumPy数组[0, 1, 2, 2, 3, 4, 2, 2, 5, 5, 6, 5, 5, 2, 2],对于每个重复序列(从每个索引开始),是否有一种快速的方法来查找该重复序列的所有匹配项并返回这些匹配项的索引?

在这里,重复序列是[2, 2][5, 5](请注意,重复的长度由用户指定,但长度相同,可以远大于2)。重复序列可以在[2,6,8,11,13]处通过以下方式找到:

def consec_repeat_starts(a, n):
    N = n-1
    m = a[:-1]==a[1:]
    return np.flatnonzero(np.convolve(m,np.ones(N, dtype=int))==N)-N+1

但是对于每种独特的重复序列类型(即[2, 2][5, 5]),我想返回类似重复的内容,然后返回重复位置的索引:

[([2, 2], [2, 6, 13]), ([5, 5], [8, 11])]

更新

此外,给定重复序列,还可以从第二个数组返回结果。因此,请在以下位置查找[2, 2][5, 5]

[2, 2, 5, 5, 1, 4, 9, 2, 5, 5, 0, 2, 2, 2]

该函数将返回:

[([2, 2], [0, 11, 12]), ([5, 5], [2, 8]))]

1 个答案:

答案 0 :(得分:0)

这是一种方法-

def group_consec(a, n):
    idx = consec_repeat_starts(a, n)
    b = a[idx]
    sidx = b.argsort()
    c = b[sidx]
    cut_idx = np.flatnonzero(np.r_[True, c[:-1]!=c[1:],True])
    idx_s = idx[sidx]
    indices = [idx_s[i:j] for (i,j) in zip(cut_idx[:-1],cut_idx[1:])]
    return c[cut_idx[:-1]], indices

# Perform lookup in another array, b
n = 2
v_a,indices_a = group_consec(a, n)
v_b,indices_b = group_consec(b, n)

idx = np.searchsorted(v_a, v_b)
idx[idx==len(v_a)] = 0
valid_mask = v_a[idx]==v_b
common_indices = [j for (i,j) in zip(valid_mask,indices_b) if i]
common_val = v_b[valid_mask]

请注意,为了简单易用,group_consec的第一个输出arg每个序列都具有唯一值。如果您需要(val, val,..)格式的文件,只需在末尾复制即可。同样,对于common_val