如何在numpy中对统计词频的简单for循环做向量化优化

解决方案 #

方案1：广播实现（写法简单，适合小规模数据） #

直接利用numpy的广播机制对两个数组做逐元素对比，再按维度求和即可，一行就能实现：

import numpy as np

unique_words = np.array(['a', 'b', 'c', 'd'])
array_to_compare = np.array(['a', 'b', 'a', 'd'])

vector_array = (array_to_compare == unique_words[:, np.newaxis]).sum(axis=1).astype(float)
# 输出：array([2., 1., 0., 1.])，和预期结果完全一致

原理是把unique_words扩展为形状为(n, 1)的二维数组，和形状为(m,)的array_to_compare广播后得到(n,m)的布尔矩阵，每一行对应当前单词和array_to_compare所有元素的对比结果，对每行求和就是当前单词的出现次数。

方案2：bincount实现（性能最优，适合大规模数据） #

如果两个数组的元素量很大，广播会占用过多内存，用np.searchsorted结合np.bincount的方案时间和空间复杂度都更低：

import numpy as np

unique_words = np.array(['a', 'b', 'c', 'd'])
array_to_compare = np.array(['a', 'b', 'a', 'd'])

# 先获取array_to_compare中每个元素在unique_words里的位置索引
sorter = np.argsort(unique_words)
pos = sorter[np.searchsorted(unique_words, array_to_compare, sorter=sorter)]
# 统计每个索引的出现次数，没有出现的默认补0，最后转成float类型
vector_array = np.bincount(pos, minlength=len(unique_words)).astype(float)
# 输出：array([2., 1., 0., 1.])

补充说明：你之前尝试np.isin没有得到预期结果，是因为该接口仅返回元素是否存在的布尔值，无法对应到每个唯一单词的位置单独计数，所以不适合这个场景。

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