Python中如何最快基于指定逻辑对两个大型列表做对应索引条件筛选

你的代码存在三处核心性能损耗：

• 每次执行shift时用列表推导式重建整个A数组，时间复杂度为O(N)，属于冗余操作
• 计算maxb和minw时两次全量遍历数组，还额外生成了无必要的中间列表，既浪费内存又增加耗时
• C是固定字符串，每次遍历都重复判断对应位置的字符，属于重复计算

前置预处理（仅执行一次） #

提前提取固定索引、替换列表为双端队列，将移位操作从O(N)降到O(1)：

from collections import deque

# 预计算B类、W类元素的下标，仅执行一次
B_IDX = [i for i, c in enumerate(C) if c == 'B']
W_IDX = [i for i, c in enumerate(C) if c == 'W']
# 将A转为双端队列，支持O(1)时间旋转
A = deque(A)

优化后的shift函数 #

def shift(A: deque, m: int, B_IDX: list, W_IDX: list) -> tuple[deque, int]:
    # 等价于原逻辑的右移一位，O(1)时间完成
    A.rotate(1)
    # 直接通过预存索引取值，无额外判断逻辑，且用生成器避免生成中间列表
    maxb = max(A[i] for i in B_IDX)
    minw = min(A[i] for i in W_IDX)
    m = max(m, maxb - minw)
    return A, m

如果B和W的索引总长度接近数组总长度，可以改成一次遍历同时计算两个值，再减少一次遍历开销：

def shift(A: deque, m: int, C: str, N: int) -> tuple[deque, int]:
    A.rotate(1)
    maxb = float('-inf')
    minw = float('inf')
    for i in range(N):
        val = A[i]
        c = C[i]
        if c == 'B' and val > maxb:
            maxb = val
        elif c == 'W' and val < minw:
            minw = val
    m = max(m, maxb - minw)
    return A, m

使用NumPy向量化操作，性能比纯Python实现高10~100倍：

import numpy as np

# 前置预处理，仅执行一次
A_np = np.array(A)
# 预生成B、W类掩码
mask_B = np.array([c == 'B' for c in C], dtype=bool)
mask_W = np.array([c == 'W' for c in C], dtype=bool)

def shift_np(A: np.ndarray, m: int, mask_B: np.ndarray, mask_W: np.ndarray) -> tuple[np.ndarray, int]:
    # 高度优化的移位操作
    A = np.roll(A, 1)
    # 向量化筛选+聚合，无Python层循环
    maxb = A[mask_B].max()
    minw = A[mask_W].min()
    m = max(m, maxb - minw)
    return A, m

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