基于np.isclose()统计浮点数组中目标浮点数的总出现次数——性能优化方案问询

你的问题核心在于原方案的时间复杂度是O(M*N)（M是keys长度，N是vals长度），当数据量达到1e6和1e7量级时，1e13次的底层操作必然导致速度极慢，再加上Python循环的额外开销，效率会更低。下面给出两种更高效的实现思路，利用NumPy的向量化运算大幅提升速度，同时避免内存溢出问题。

方案1：分块向量化处理（通用场景） #

将keys分成小批量，每次用NumPy的向量化操作批量计算与vals的匹配情况，替代逐个key的Python循环。这种方法充分利用了NumPy的C级运算效率，同时通过分块控制内存占用。

import numpy as np

def count_float_keys_fast(keys, vals):
    keys_arr = np.asarray(keys)
    vals_arr = np.asarray(vals)
    total_count = 0
    # 分块大小可根据你的内存情况调整，比如1000/500/100
    chunk_size = 1000
    
    for i in range(0, len(keys_arr), chunk_size):
        # 取出当前批次的keys
        key_chunk = keys_arr[i:i+chunk_size]
        # 广播vals到批次维度，判断每个val是否匹配批次中的任意key
        matches = np.isclose(vals_arr[:, None], key_chunk).any(axis=1)
        # 累加当前批次的匹配总数
        total_count += matches.sum()
    
    return total_count

方案2：先去重keys再统计（keys含大量重复时更优） #

如果你的keys列表中有大量重复值，可以先对keys去重并统计每个唯一key的出现次数，再计算每个唯一key在vals中的匹配数，最后加权求和。这种方法能减少需要处理的key数量，进一步提升效率。

import numpy as np

def count_float_keys_fast(keys, vals):
    keys_arr = np.asarray(keys)
    vals_arr = np.asarray(vals)
    # 去重并统计每个key的出现次数
    unique_keys, key_counts = np.unique(keys_arr, return_counts=True)
    total_count = 0
    chunk_size = 1000
    
    for i in range(0, len(unique_keys), chunk_size):
        chunk_keys = unique_keys[i:i+chunk_size]
        chunk_counts = key_counts[i:i+chunk_size]
        # 统计每个唯一key在vals中的匹配数
        matches_per_key = np.isclose(vals_arr[:, None], chunk_keys).sum(axis=0)
        # 加权累加：匹配数 * 该key的出现次数
        total_count += (matches_per_key * chunk_counts).sum()
    
    return total_count

为什么这两种方案更快？ #

• 替代Python循环：NumPy的向量化操作是在C层面执行的，比Python循环快几个数量级。
• 减少重复计算：方案2通过去重keys，避免了对同一个key重复计算与vals的匹配。
• 内存可控：分块处理避免了直接创建1e7*1e6的巨型数组，防止内存溢出。

测试示例 #

用你给出的测试数据验证：

keys = [0.999999, 1.999999]
vals = [1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 1.0, 2.0]
print(count_float_keys_fast(keys, vals))  # 输出5，符合预期

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