如何Cython化NumPy数组以加速循环及isSubsetquick函数性能？

一、核心优化思路 #

用Cython给NumPy数组指定静态类型，消除Python动态类型的开销；同时把核心函数改写成Cython静态类型函数，减少函数调用和类型检查的损耗，最终提升外部大循环的运行速度。

二、Cython化isSubsetquick函数 #

直接给函数参数和内部变量指定静态类型，让Cython直接编译成C级代码：

import numpy as np
cimport numpy as np

# 根据实际数据类型调整，这里用int类型示例
cdef int isSubsetquick(np.ndarray[np.int_t, ndim=3] arr_3d, np.ndarray[np.int_t, ndim=1] arr_1d, int limit_elements):
    cdef np.ndarray[np.bool_t, ndim=3] common_elements
    cdef np.ndarray[np.int_t, ndim=2] quantity_common_elements
    cdef int counter

    # 直接调用NumPy的底层接口，跳过Python层包装
    common_elements = np.isin(arr_3d, arr_1d[:limit_elements])
    quantity_common_elements = np.sum(common_elements, axis=-1, dtype=np.int_)
    counter = np.count_nonzero(quantity_common_elements == limit_elements)

    return counter

三、优化外部大型循环 #

把循环变量、数组都指定为Cython静态类型，同时移除循环内的IO操作（比如print），避免拖慢速度：

def main():
    # 提前指定数组类型，避免Cython自动推导的开销
    cdef np.ndarray[np.int_t, ndim=3] array_3 = np.array(
        [[[1, 2, 3, 7], [4, 5, 6, 8]], [[7, 8, 9, 4], [10, 11, 12, 5]]],
        dtype=np.int_
    )
    cdef np.ndarray[np.int_t, ndim=1] array_1 = np.array([1, 2, 6], dtype=np.int_)
    
    cdef int limit = 3
    cdef int size = 100000
    cdef int result
    cdef int x

    # 循环内只做核心计算，print移出循环减少IO开销
    for x in range(size):
        result = isSubsetquick(array_3, array_1, limit)
    print(result)

四、额外提速技巧 #

• 统一数组类型：所有NumPy数组都明确指定dtype（比如np.int_、np.float64），避免Cython处理隐式类型转换
• 提前预处理：如果limit在循环中不变，提前把arr_1d[:limit]切片好传入函数，避免每次循环重复切片
• 编译优化：编译Cython文件时，在setup.py里添加编译参数extra_compile_args=['-O3', '-march=native']，让编译器生成最优的机器码

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