轨迹合并学习笔记

import numpy as np

def find_contiguous_indices(mask_all, max_gap=3):
    # 初始化变量
    segments = []
    n = len(mask_all)
    i = 0

    # 先找到所有连续的 True 段
    while i < n:
        if mask_all[i]:
            start = i
            while i < n and mask_all[i]:
                i += 1
            end = i - 1
            segments.append((start, end))
        else:
            i += 1

    # 合并间隔 < max_gap 的段
    if not segments:
        return np.array([], dtype=int)  # 如果没有 True，返回空数组

    merged_segments = [segments[0]]
    for current in segments[1:]:
        last = merged_segments[-1]
        if current[0] - last[1] - 1 < max_gap:
            merged_segments[-1] = (last[0], current[1])  # 合并
        else:
            merged_segments.append(current)

    # 生成所有 True 的索引列表
    indices = []
    for start, end in merged_segments:
        indices.append([i for i in range(start, end + 1)])  # 添加该段所有索引

    return indices

# 示例 mask_all
mask_all = np.array([
    True, True, False, False, True, False, True, True, True,
    False, False, False, False, True, True, False, True
], dtype=bool)

print("原始 mask_all:", mask_all)
indices = find_contiguous_indices(mask_all, max_gap=3)

for index in indices:
    print(index)