大量整数数组匹配计数的高效算法优化方案问询

现有场景是统计数百万个整数数组（每个数组含固定数量元素）与参考整数数组的匹配元素数量，当前已实现基于Cython的方案，利用指示器数组快速判断匹配。由于实际场景中所有整数数值不大，且唯一整数数量仅数千个（远小于数组总元素数），可针对性优化如下：

现有实现代码 #

1. Cython核心实现 (match_ints.pyx) #

cimport cython
from libc.stdlib cimport calloc, free

import numpy as np
cimport numpy as np

np.import_array()


@cython.wraparound(False)
@cython.boundscheck(False)
@cython.initializedcheck(False)
cdef void count_matches(int[:, ::1] target_arrays, int[::1] ref_array, int[::1] num_matches):

    cdef:
        Py_ssize_t i, j
        Py_ssize_t n = target_arrays.shape[0]
        Py_ssize_t c = target_arrays.shape[1]
        Py_ssize_t nf = ref_array.shape[0]
        Py_ssize_t m = ref_array[nf - 1] + 5
        int * ind = <int *> calloc(m, sizeof(int))
        int k, g

    for i in range(nf):
        ind[ref_array[i]] = 1

    for i in range(n):
        k = 0
        for j in range(c):
            g = target_arrays[i, j]
            if g < m and ind[g] == 1:
                k += 1
        num_matches[i] = k

    free(ind)


cpdef count_num_matches(int[:, ::1] target_arrays, int[::1] ref_array):

    cdef:
        Py_ssize_t n = target_arrays.shape[0]
        int[::1] num_matches = np.zeros(n, dtype=np.int32)

    count_matches(target_arrays, ref_array, num_matches)

    return np.asarray(num_matches)

2. 测试代码 (test_main_counts.py) #

from match_ints import count_num_matches
import numpy as np


def count_num_matches_main():
    x = np.random.randint(50, 6000, size=(1000000, 40), dtype=np.int32)
    ref_x = np.random.randint(100, 2500, size=800, dtype=np.int32)

    ref_x.sort()

    return count_num_matches(x, ref_x)


if __name__ == "__main__":
     nums = count_num_matches_main()
     print(nums[:10])

3. 构建脚本 (setup.py) #

from setuptools import setup
from Cython.Build import cythonize
import numpy as np


setup(
    ext_modules=cythonize(
        "match_ints.pyx",
        compiler_directives={
            "language_level": "3",
        }
    ),
    include_dirs=[
        np.get_include()
    ]
)

基于唯一整数特性的优化方案 #

针对唯一整数数量少的特点，可从减少重复判断和批量向量化处理两个方向优化：

1. 预唯一值匹配映射 #

• 先提取所有目标数组的唯一整数集合，仅对这些唯一值判断是否存在于参考数组中，生成匹配映射表（键为整数，值为1/0表示是否匹配）。
• 遍历目标数组时，直接通过映射表查询元素是否匹配，避免对重复元素多次判断。

2. Numpy向量化批量实现 #

利用Numpy的高效向量化操作，结合唯一值编码减少计算量：

def optimized_count(x, ref_x):
    ref_set = set(ref_x)
    # 提取目标数组唯一值并生成编码映射
    unique_vals, encoded_indices = np.unique(x, return_inverse=True)
    # 生成唯一值的匹配标记
    match_mask = np.array([1 if val in ref_set else 0 for val in unique_vals], dtype=np.int32)
    # 将编码还原为原数组形状，快速求和每行匹配数
    encoded_x = encoded_indices.reshape(x.shape)
    counts = np.sum(match_mask[encoded_x], axis=1)
    return counts

该方案中每个唯一值仅做一次匹配判断，后续通过索引映射完成批量计算，内存占用和计算时间均显著降低。

3. Cython层面的针对性优化 #

在原有Cython代码基础上，结合唯一值特性改进：

• 用哈希集合替代指示器数组：若参考数组的最大值较大但唯一值少，哈希集合的内存占用更低，判断匹配的时间复杂度仍为O(1)。
• 预生成目标唯一值的匹配标记：提前对目标数组的唯一值做匹配判断，遍历目标数组时直接读取标记，减少条件分支判断。

示例修改后的Cython核心逻辑：

# 替换原指示器数组逻辑为哈希集合
cdef:
    set ref_set = set(ref_array)
    # 预生成目标唯一值的匹配字典（可提前在Python层处理后传入）
    dict match_map = {val: 1 if val in ref_set else 0 for val in np.unique(target_arrays)}

# 遍历目标数组时直接查字典
for i in range(n):
    k = 0
    for j in range(c):
        g = target_arrays[i, j]
        k += match_map.get(g, 0)
    num_matches[i] = k

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