Python中高效计算数组元素到左侧最近1的步数方法

问题描述 #

给定数组foo，需要生成数组bar，其中bar[i]表示从索引i向左找到最近的1所需的步数：

import numpy as np

foo = np.array([1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0])
# 目标输出
bar = np.array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3])

纯NumPy高效解法 #

利用np.where定位所有1的位置，结合np.searchsorted快速匹配每个索引对应的最近左侧1的位置，无需循环：

ones_idx = np.where(foo == 1)[0]
# 找到每个索引在ones_idx中的插入位置，取前一个1的索引
pos = np.searchsorted(ones_idx, np.arange(len(foo)), side='right') - 1
bar = np.arange(len(foo)) - ones_idx[pos]

原理说明 #

1. np.where(foo == 1)[0]获取所有值为1的索引；
2. np.searchsorted(..., side='right')返回每个索引在ones_idx中右侧插入点，减1后得到最近的左侧1的索引；
3. 用当前索引减去该1的索引，即为所需步数。

Numba优化循环解法 #

普通Python循环速度极慢，但通过Numba即时编译（JIT）可将循环转换为机器码，性能与纯NumPy方案持平：

from numba import jit

@jit(nopython=True)
def compute_steps_numba(arr):
    n = len(arr)
    result = np.zeros(n, dtype=np.int64)
    last_one_pos = 0
    for i in range(n):
        if arr[i] == 1:
            last_one_pos = i
            result[i] = 0
        else:
            result[i] = i - last_one_pos
    return result

bar_numba = compute_steps_numba(foo)

性能对比 #

• Pandas Series方案：比纯NumPy方案慢约7倍；
• 普通Python循环：速度极慢，不推荐；
• Numba编译后的循环：与纯NumPy方案速度相当。

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