在NumPy数组中获取最小长度的相同条目序列范围

Question

考虑一个只包含-1或1的数组。如何获取所有仅包含1且长度不小于t (例如t = 3)的切片范围。

例如：

>>>a=np.array([-1,-1,1,1,1,1,1,-1,1,-1,-1,1,1,1,1], dtype=int)
>>> a
array([-1, -1,  1,  1,  1,  1,  1, -1,  1, -1, -1,  1,  1,  1,  1])

那么，当t=3时，期望的输出应该是[(2,7),(11,15)]。

- corinna

Answer 1

使用np.diff和np.where的一种方法 -

# Append with `-1s` at either ends and get the differentiation
dfa = np.diff(np.hstack((-1,a,-1)))

# Get the positions of starts and stops of 1s in `a`
starts = np.where(dfa==2)[0]
stops = np.where(dfa==-2)[0]

# Get valid mask for pairs from starts and stops being of at least 3 in length
valid_mask = (stops - starts) >= 3

# Finally collect the valid pairs as the output
out = np.column_stack((starts,stops))[valid_mask].tolist()

- Divakar

---

@corinna 很有趣，所以完全没有问题！ - Divakar

Answer 2

虽然我不是很熟悉numpy，但使用简单的函数会更好吧？

def slices(a, t):
    start = None
    i = 0 # index into array
    slices = [] 
    for val in a:
        if a[i] == 1: # start of sequence
            if start is None:
                start = i
        else: # -1 end of sequence
            if start is not None:
                if i - start >= t: # check sequence for minimum size
                    slices.append((start, i))
                start = None
        i += 1

    # if sequence of 1's doesn't end with -1 within array
    if start is not None:
        if i - start >= t:
            slices.append((start, i))

   return slices

- Radek Luner