问将索引重新对齐到已更改的python集合

EN

import random

def align_index(wanted_idx, mask):
    """
    Function to align a set of indexes to a collection after deletions,
    indicated with a mask
    Arguments:
    wanted_idx: List of desired integer indexes prior to deletion
    mask: Binary mask, where 1's indicate elements that survive deletion
    Returns:
    List of integer indexes to (surviving) desired elements, post-deletion
    """
    # rebuild indexes: remove dangling
    new_idx = [idx for (i, idx) in enumerate(wanted_idx) if mask[idx]]
    # mark deleted
    not_mask = [int(not m) for m in mask]
    # cumsum deleted regions
    realigned_idx = [k-sum(not_mask[:k+1]) for k in new_idx]
    return realigned_idx

# data
data = [random.randint(0,500) for _ in range(1000)]
rng = list(range(len(data)))

for _ in range(1000):
    # random data deletion / request
    wanted_idx = random.sample(rng, random.randint(5,100))
    del_index = random.sample(rng, random.randint(5, 100))

    # apply deletion
    mask = [int(i not in del_index) for i in range(len(data))]
    filtered_data = [data[i] for (i, m) in enumerate(mask) if m]

    realigned_index = align_index(wanted_idx, mask)

    # verify
    new_idx = [idx for (i, idx) in enumerate(wanted_idx) if mask[idx]]
    l1 = [data[k] for k in new_idx]
    l2 = [filtered_data[k] for k in realigned_index]
    assert l1 == l2