如何在NumPy中快速查找末尾符合条件的连续元素数量？

嘿，针对你这个在大NumPy数组里统计末尾连续1数量的需求，我整理了几个高效的方案——完全不用循环，还能避免全数组遍历，特别适合超大数组的场景：

这个方法是效率最高的，通过不断缩小范围来定位最长的末尾连续1长度，哪怕数组有几百万个元素，也只需要几十次检查，完全不用遍历整个数组：

import numpy as np

x = np.array([1, 2, 3, 4, 5, 1, 1, 1])

left, right = 0, len(x)
max_count = 0

while left <= right:
    mid = (left + right) // 2
    # 检查末尾mid个元素是否全为1
    if mid == 0:
        current_valid = True
    else:
        current_valid = (x[-mid:] == 1).all()
    
    if current_valid:
        max_count = mid
        left = mid + 1  # 尝试更长的连续长度
    else:
        right = mid - 1  # 缩短长度尝试

print(max_count)  # 输出: 3

如果你的数组不是特别极端的大，这个方法代码更简洁，利用NumPy的向量化操作实现：

import numpy as np

x = np.array([1, 2, 3, 4, 5, 1, 1, 1])

# 反转数组，标记出所有不是1的元素
rev_non_one = x[::-1] != 1
if rev_non_one.any():
    # 找到反转后第一个非1元素的索引，就是原数组末尾连续1的数量
    count = rev_non_one.argmax()
else:
    # 处理数组全是1的边界情况
    count = len(x)

print(count)  # 输出: 3

注意：argmax会遍历整个反转后的数组，所以如果数组极大且末尾几乎全是1，这个方法的效率不如二分查找胜在代码简洁易读。

通过找到最后一个不是1的元素的索引，直接计算末尾连续1的数量：

import numpy as np

x = np.array([1, 2, 3, 4, 5, 1, 1, 1])

non_one_indices = np.where(x != 1)[0]
if not len(non_one_indices):
    # 数组全为1的情况
    count = len(x)
else:
    last_non_one_pos = non_one_indices[-1]
    count = len(x) - last_non_one_pos - 1

print(count)  # 输出: 3

这个方法会收集所有非1元素的索引，适合数组中非1元素不多的场景，但极端大数组下还是不如二分查找高效。

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