python - 数组的最长前缀等于一个值

Question

我有 NumPy 字符串数组，表示序列的分段。 B 是感兴趣段的开始，I 是它的延续，O 在任何段之外。例如，在下面的数组中，有三个感兴趣的段:

>>> y
array(['B', 'I', 'I', 'O', 'B', 'I', 'O', 'O', 'B', 'O'], 
      dtype='|S1')

我可以找到以 np.where(y == "B")[0] 轻松开始的段。但现在我还试图找到段的长度，即等于 I 的最长前缀的长度。我可以使用 itertools.takewhile 这样做:

>>> from itertools import takewhile
>>> lengths = [1 + sum(1 for _ in takewhile(lambda x: x == "I", y[start + 1:]))
...            for start in np.where(y == "B")[0]]
>>> lengths
[3, 2, 1]

老实说，它工作得很好，但是有没有矢量化的方法来实现这个目标？

Answer 1

段以 'B' 开头:

starts = np.where(y == 'B')[0]

一个段结束于 'B' 或 'I' 之后是 'I' 以外的东西，或者是序列结束:

ends = np.where(((y == 'B') | (y == 'I')) & np.r_[y[1:] != 'I', len(y)])[0]

这给出了段长度:

(ends - starts) + 1
array([3, 2, 1])

---

编辑:这里有一个更简单的方法:在末尾插入一个虚构的 B，然后取(真实的或虚构的)B 位置的差异，不包括O:

np.diff(np.where(np.r_[y[y != 'O'], ['B']] == 'B')[0])
array([3, 2, 1])