Python：实现双列表列表结构的线性__getitem__方法

需求说明 #

• 类需同时存储两个等长的列表列表：instructions和annotations（子列表长度可不同）
• 通过__getitem__实现线性下标访问，遍历顺序为交替访问两列表的对应子列表：先取instructions[0]的所有元素，再取annotations[0]的所有元素，接着instructions[1]、annotations[1]，以此类推
• 支持正负索引，示例：y[9]返回9、y[-4]返回13
• 实现需高效，禁止循环遍历（适配大数据量场景）

实现思路 #

1. 预处理前缀和数组：在类初始化时，预先计算三个前缀和数组：
   • inst_prefix：记录instructions中前i个子列表的总元素数
   • anno_prefix：记录annotations中前i个子列表的总元素数
   • combined_prefix：记录前i对（instructions[i]+annotations[i]）的总元素数
2. 索引转换：将负索引转换为正索引，避免边界判断混乱
3. 二分查找定位：利用二分查找快速定位目标索引属于哪一对子列表，再计算其在对应子列表中的偏移量，直接返回元素

代码实现 #

import bisect

class NewClass:
    def __init__(self, instructions, annotations):
        # 校验两个列表长度相等
        if len(instructions) != len(annotations):
            raise ValueError("instructions and annotations must be of the same length")
        
        self.instructions = instructions
        self.annotations = annotations
        
        # 预处理前缀和数组
        self.inst_prefix = [0]
        self.anno_prefix = [0]
        self.combined_prefix = [0]
        
        for inst_sub, anno_sub in zip(instructions, annotations):
            self.inst_prefix.append(self.inst_prefix[-1] + len(inst_sub))
            self.anno_prefix.append(self.anno_prefix[-1] + len(anno_sub))
            self.combined_prefix.append(self.combined_prefix[-1] + len(inst_sub) + len(anno_sub))
    
    def __len__(self):
        return self.combined_prefix[-1]
    
    def __getitem__(self, idx):
        # 处理负索引
        if idx < 0:
            idx += len(self)
        if idx < 0 or idx >= len(self):
            raise IndexError("NewClass index out of range")
        
        # 二分查找找到对应的子列表对
        k = bisect.bisect_right(self.combined_prefix, idx) - 1
        offset_in_pair = idx - self.combined_prefix[k]
        
        # 判断属于instructions还是annotations的子列表
        inst_sub_len = len(self.instructions[k])
        if offset_in_pair < inst_sub_len:
            return self.instructions[k][offset_in_pair]
        else:
            return self.annotations[k][offset_in_pair - inst_sub_len]

# 示例验证
if __name__ == "__main__":
    # 构造符合示例的测试数据
    instructions = [[0,1,2,3,4], [8,9,10,11]]
    annotations = [[5,6,7], [12,13,14,15,16]]
    y = NewClass(instructions, annotations)
    print(y[9])   # 输出9
    print(y[-4])  # 输出13

代码说明 #

• 初始化阶段：一次性计算三个前缀和数组，后续访问无需重复计算，保证效率
• 索引处理：负索引转换为正索引后，通过bisect_right快速定位所在的子列表对，时间复杂度为O(log n)（n为子列表对的数量）
• 无循环实现：整个访问过程没有遍历元素，仅通过前缀和与二分查找完成定位，完全适配大数据量场景

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