如何用Fortran和f2py实现比np.where更快的数组值查找？

我来帮你拆解下这个性能瓶颈问题，以及给出具体的优化方向——你遇到的np.where性能瓶颈，还有用Fortran+f2py没提效的情况，我之前帮不少开发者解决过类似问题，核心要么是Python端的循环/交互开销没控制好，要么是Fortran与numpy的数据交互没利用好语言特性。下面分Python和Fortran两个方向给你落地的优化方案：

Python端优化方案 #

1. 批量处理数组，减少循环开销 #

如果你的数百个数组是独立的，优先把它们堆叠成更高维度的数组（比如形状为(n_arrays, rows, cols)），一次性用np.where处理，避免循环遍历每个数组带来的Python层交互开销：

import numpy as np

# 假设你有数百个形状为(1000,1000)的整数数组，存在列表arr_list中
stacked_arr = np.stack(arr_list)
target_val = 42

# 一次性获取所有数组中目标值的索引
indices = np.where(stacked_arr == target_val)
# indices是三元组：(数组索引, 行索引, 列索引)，可拆分对应原数组

这种方式把多次np.where调用合并成一次，大幅减少Python与底层C的交互次数，性能提升非常明显。

2. 用更轻量的索引提取方法 #

np.where本质是生成布尔掩码再提取索引，你可以直接用np.nonzero配合数组比较，少一层封装的情况下性能会略优：

# 单个数组的情况
arr = np.random.randint(0, 100, size=(1000,1000))
indices = np.nonzero(arr == target_val)

如果目标值在数组中出现频率极低，np.argwhere也是可选方案（返回(N, ndim)形状的索引数组，按需调整即可）。

3. 利用Numba即时编译 #

Numba能把Python函数编译成机器码，彻底绕过Python解释器的循环开销，尤其适合逐个处理数组的场景：

from numba import jit

@jit(nopython=True)
def find_indices_numba(arr, target):
    rows = []
    cols = []
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i,j] == target:
                rows.append(i)
                cols.append(j)
    return np.array(rows), np.array(cols)

# 调用示例
rows, cols = find_indices_numba(arr, target_val)

nopython模式下的性能接近编译型语言，而且不需要切换到Fortran，学习成本极低。

Fortran端优化方案（解决你之前性能没提升的问题） #

你之前用Fortran没提效，大概率是没处理好numpy与Fortran的内存布局差异，或者代码没利用Fortran的向量化/缓存优势。下面是优化后的代码和调用注意事项：

1. 优化的Fortran代码（适配内存布局） #

Fortran是列优先内存布局，和numpy默认的行优先不同，直接遍历列能提升缓存命中率：

subroutine find_indices_fortran(arr, nrows, ncols, target, n_indices, rows, cols)
    implicit none
    integer, intent(in) :: nrows, ncols, target
    integer, intent(in) :: arr(nrows, ncols)
    integer, intent(out) :: n_indices
    integer, intent(out) :: rows(nrows*ncols), cols(nrows*ncols)
    integer :: i, j, count

    count = 0
    ! 按列遍历，符合Fortran内存布局，减少缓存 miss
    do j = 1, ncols
        do i = 1, nrows
            if (arr(i,j) == target) then
                count = count + 1
                rows(count) = i - 1  ! 转换为Python的0-based索引
                cols(count) = j - 1
            end if
        end do
    end do
    n_indices = count
end subroutine find_indices_fortran

2. 编译与调用的关键细节 #

• 编译时开启最高优化级别：

  f2py -c -O3 -m find_indices find_indices.f90
  

• Python调用时，确保numpy数组是Fortran连续的，避免不必要的内存拷贝：

  import numpy as np
  import find_indices
  
  # 创建Fortran连续的数组，避免转换开销
  arr = np.random.randint(0, 100, size=(1000,1000), order='F')
  target_val = 42
  max_indices = arr.size
  rows = np.zeros(max_indices, dtype=np.int32)
  cols = np.zeros(max_indices, dtype=np.int32)
  
  n_indices = find_indices.find_indices_fortran(arr, arr.shape[0], arr.shape[1], target_val, max_indices, rows, cols)
  # 提取有效索引
  valid_rows = rows[:n_indices]
  valid_cols = cols[:n_indices]
  

  这一步的order='F'是核心，很多人用f2py性能上不去，就是因为忽略了内存布局导致的隐式数据拷贝。

3. 进一步优化：向量化Fortran实现 #

Fortran编译器对向量化支持极佳，改用向量化写法能让编译器自动做循环展开、缓存优化：

subroutine find_indices_vectorized(arr, nrows, ncols, target, n_indices, rows, cols)
    implicit none
    integer, intent(in) :: nrows, ncols, target
    integer, intent(in) :: arr(nrows, ncols)
    integer, intent(out) :: n_indices
    integer, intent(out) :: rows(nrows*ncols), cols(nrows*ncols)
    logical :: mask(nrows, ncols)

    mask = arr == target
    ! 用pack函数批量提取符合条件的索引
    rows = pack([(i, i=1,nrows)], mask) - 1
    cols = pack([(j, j=1,ncols)], mask) - 1
    n_indices = count(mask)
end subroutine find_indices_vectorized

总结 #

• 不想切换语言的话，Numba是最便捷的选择，性能接近编译型语言且代码风格和Python一致；
• 坚持用Fortran的话，一定要处理好内存布局（避免数据拷贝），并开启编译优化；
• 批量处理数组能大幅减少Python层的循环开销，优先考虑堆叠数组后一次性处理。

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