numpy 在包含0和1的列表中查找1的起点和终点

iyfjxgzm  于 2022-11-24  发布在  其他
关注(0)|答案(4)|浏览(148)

我有一个列表,我想找到值1的起始和结束索引。

labels=[0,0,0,1,1,1,0,0,1,1]

1的索引为[3,5][8,9]
有没有什么有效的方法可以做到这一点?我试过numpy index(),但是它对我不起作用。使用它,我可以找到第一个或最后一个1,但是不能找到中间的1。这就是我所尝试的。

[labels.index(1),len(labels)-labels[::-1].index(1)-1]

这样得到的是[3,9],但我希望得到连续1的索引,即[3,5][8,9]

bf1o4zei

bf1o4zei1#

一个选项是使用diff获取每个拉伸的开始,并在填充后与0进行比较,以确保标签以0开始和结束:

labels = [0,0,0,1,1,1,0,0,1,1]

out = np.where(np.diff(np.pad(labels, 1))!=0)[0].reshape(-1,2)-[0,1]

输出:

array([[3, 5],
       [8, 9]])

作为列表:

out.tolist()
# [[3, 5], [8, 9]]
46qrfjad

46qrfjad2#

@mozway建议的方法可能更适合大型列表,但是,对于这样一个小数据集,这种方法要快得多(如果这很重要的话):

labels = [0, 0, 0, 1, 1, 1, 0, 0, 1, 1]

def func(labels):
    result = []
    se = -1
    for i, e in enumerate(labels):
        if e == 1:
            if se < 0:
                se = i
        else:
            if se >= 0:
                result.append([se, i-1])
                se = -1
    if se >= 0:
        result.append([se, len(labels)-1])
    return result

print(func(labels))

输出:

[[3, 5], [8, 9]]
5hcedyr0

5hcedyr03#

一种方法是使用itertools compress,groupby和count。“获取索引1,按连续索引分组,并找到组中的第一个和最后一个索引”。

from itertools import groupby, compress, count

serie = [0,0,0,1,1,1,0,0,1,1,0,1,0]      

stream = compress(*zip(*enumerate(serie)))    # indices of 1

consecutives = ([*streak] for _, streak in groupby(stream, lambda n,  c=count(): n - next(c)))  # streaks of consecutive numbers

indices = [[el[0], el[-1]] for el in consecutives]

-> [[3, 5], [8, 9], [11, 11]]
vu8f3i0k

vu8f3i0k4#

对于numpy和python初学者:

labels=[0,0,0,1,1,1,0,0,1,1,0,1]
indices=np.array(np.nonzero(labels))[0]
count,result = 0 , []
for idx in range(len(indices)-1):
    if indices[idx]+1==indices[idx+1]:
        count+=1
        if idx==len(indices)-2:
            result.append([indices[idx],indices[idx]+1])
    else:
        result.append([indices[idx]-count,indices[idx]])
        count=0
if indices[-1]-indices[-2]>1:
    result.append([indices[-1],indices[-1]])
print(result)

Output : [[3, 5], [8, 9], [11, 11]]

相关问题