基于偏移量的numpy‘不规则’数组重排高效实现方案问询

现有一个1D字节型数组（如字符数组），以及一组定义连续切片（对应“单词”）的偏移量数组，示例如下：

import numpy as np
data = np.array(['h', 'e', 'l', 'l', 'o', 't', 'h', 'e', 'r', 'e', 'y', 'o', 'u', '!'])
offsets = np.array([0, 5, 10, 13, 14]) # 对应切片 [0:5], [5:10], [10:13], [13:14]

需要实现高效方法按指定顺序重排这些切片，支持重复或删除操作，示例：

1. 重排且长度不变的情况：

new_word_order = np.array([2, 1, 0, 3])
# 期望输出：
new_data == np.array(['y', 'o', 'u', 't', 'h', 'e', 'r', 'e', 'h', 'e', 'l', 'l', 'o', '!'])
new_offsets == np.array([0, 3, 8, 13, 14])

2. 重复/删除切片的情况：

new_word_order = np.array([2, 3, 3, 2, 3])
# 期望输出：
new_data == np.array(['y', 'o', 'u', '!', '!', 'y', 'o', 'u', '!'])
new_offsets = np.array([0, 3, 4, 5, 8, 9])

原实现基于列表推导式循环调用np.arange，但百万级切片场景下效率不足，希望找到内置多切片访问方法，或借助np.r_、np.s_、np.slice等工具优化。

核心思路：避免循环生成np.arange，直接构造索引数组 #

原方案中循环调用np.arange会产生大量小数组，在百万级切片场景下开销极大。可以通过计算每个切片的长度，构造重复次数数组，再用np.repeat和np.cumsum生成索引，全程用NumPy矢量运算替代Python循环。

具体实现代码 #

import numpy as np

def reorder_slices(data, offsets, new_word_order):
    # 计算原始切片的起始、结束位置和长度
    starts = offsets[:-1]
    ends = offsets[1:]
    lengths = ends - starts
    
    # 根据新顺序提取对应的起始位置和长度
    reordered_starts = starts[new_word_order]
    reordered_lengths = lengths[new_word_order]
    
    # 过滤长度为0的无效切片
    valid_mask = reordered_lengths > 0
    if not np.any(valid_mask):
        return np.array([]), np.array([0])
    
    filtered_starts = reordered_starts[valid_mask]
    filtered_lengths = reordered_lengths[valid_mask]
    
    # 构造完整索引数组：每个切片的起始位置加上0到length-1的序列
    idx = np.repeat(filtered_starts, filtered_lengths) + \
          np.arange(filtered_lengths.sum()) - \
          np.repeat(np.cumsum(filtered_lengths) - filtered_lengths, filtered_lengths)
    
    # 生成重排后的data数组
    new_data = data[idx]
    
    # 生成新的偏移量数组
    cum_lengths = np.cumsum(filtered_lengths)
    new_offsets = np.insert(cum_lengths, 0, 0)
    
    return new_data, new_offsets

方案优势 #

1. 性能高效：全程矢量运算，无Python级循环，百万级切片场景下比原方案性能提升数倍。
2. 内存友好：直接构造完整索引数组，避免生成大量中间小数组。
3. 兼容性强：天然支持切片重复、删除，以及无效切片过滤。

测试验证 #

用示例数据验证：

data = np.array(['h', 'e', 'l', 'l', 'o', 't', 'h', 'e', 'r', 'e', 'y', 'o', 'u', '!'])
offsets = np.array([0, 5, 10, 13, 14])

# 测试用例1
new_word_order1 = np.array([2, 1, 0, 3])
new_data1, new_offsets1 = reorder_slices(data, offsets, new_word_order1)
print(np.array_equal(new_data1, np.array(['y', 'o', 'u', 't', 'h', 'e', 'r', 'e', 'h', 'e', 'l', 'l', 'o', '!'])))  # 输出True
print(np.array_equal(new_offsets1, np.array([0, 3, 8, 13, 14])))  # 输出True

# 测试用例2
new_word_order2 = np.array([2, 3, 3, 2, 3])
new_data2, new_offsets2 = reorder_slices(data, offsets, new_word_order2)
print(np.array_equal(new_data2, np.array(['y', 'o', 'u', '!', '!', 'y', 'o', 'u', '!'])))  # 输出True
print(np.array_equal(new_offsets2, np.array([0, 3, 4, 5, 8, 9])))  # 输出True

补充：np.r_简化写法（小切片场景适用） #

如果切片数量不多，也可以用np.r_快速拼接切片对象，但此方法在百万级切片场景下，生成大量slice对象的开销不可忽视，仅适合小规模场景：

def reorder_slices_r(data, offsets, new_word_order):
    slices = [slice(offsets[i], offsets[i+1]) for i in range(len(offsets)-1)]
    reordered_slices = [slices[i] for i in new_word_order if slices[i].start != slices[i].stop]
    idx = np.r_[tuple(reordered_slices)]
    new_data = data[idx]
    # 计算新偏移量
    lengths = [s.stop - s.start for s in reordered_slices]
    new_offsets = np.insert(np.cumsum(lengths), 0, 0)
    return new_data, new_offsets

内容的提问来源于stack exchange，提问作者user2667066