majiayu000/claude-skill-registry-data/adjacent-span-merging

def merge_adjacent_spans(predictions, max_gap=1, merge_classes=None):
    """Merge nearby same-class spans within max_gap words.
    Args:
        predictions: list of (doc_id, class, word_indices_str)
        max_gap: max word-index gap to merge across
        merge_classes: set of classes to merge (None = all)
    Returns:
        list of merged (doc_id, class, word_indices_str)
    """
    from collections import defaultdict
    by_doc_class = defaultdict(list)
    for doc_id, cls, pred_str in predictions:
        words = sorted(int(x) for x in pred_str.split())
        by_doc_class[(doc_id, cls)].append(words)
    merged = []
    for (doc_id, cls), spans in by_doc_class.items():
        if merge_classes and cls not in merge_classes:
            for s in spans:
                merged.append((doc_id, cls, " ".join(map(str, s))))
            continue
        # Sort spans by start position
        all_words = sorted(set(w for s in spans for w in s))
        # Split into groups by gap
        groups = [[all_words[0]]]
        for w in all_words[1:]:
            if w - groups[-1][-1] <= max_gap:
                groups[-1].append(w)
            else:
                groups.append([w])
        for g in groups:
            merged.append((doc_id, cls, " ".join(map(str, g))))
    return merged