这个问题实际上很有趣，因为它是计算时间和内存消耗之间权衡的完美示例。
从算法的角度来看，找到唯一元素并最终只计算唯一元素，可以通过两种方式实现：
-
两次（或更多）通过方法：
- 找出所有独特的元素
- 找出独特元素的位置
- 计算唯一元素的函数
- 将所有计算出来的唯一元素放到正确的位置
-
单程方法：
- 随时随地计算元素并缓存结果
- 如果一个元素在缓存中，则从那里获取它
算法复杂度取决于输入 N 的大小和唯一元素的数量 U。后者也可以使用独特元素的r = U / N 比率来形式化。
more-passes 方法理论上较慢。但是，它们对于小型N 和U 具有相当的竞争力。
single-pass 方法理论上更快，但这也很大程度上取决于 缓存 方法以及它们的执行方式取决于U。
当然，无论渐近行为多么重要，实际时间确实取决于恒定的计算时间因素。
与这个问题最相关的是func() 计算时间。
方法
可以比较多种方法：
-
没有缓存
-
pure() 这将是基本函数，可能已经被矢量化了
-
np.vectorized() 这将是 NumPy 标准矢量化装饰器
-
more-passes 方法
-
np_unique()：使用np.unique() 找到唯一值并使用索引（来自np.unique() 输出）来构造结果（基本上相当于来自here 的vectorize_pure()）
-
pd_unique()：使用pd.unique() 找到唯一值并使用索引（通过np.searchsorted()）构造结果（基本上相当于vectorize_with_pandas() 来自here）
-
set_unique()：使用简单的set() 找到唯一值并使用索引（通过np.searchsorted()）构造结果
-
set_unique_msk()：使用简单的set()（如set_unique()）找到唯一值，并使用循环和屏蔽来构造结果（而不是索引）
-
nb_unique()：使用显式循环和numba JIT 加速找到唯一值及其索引
-
cy_unique()：唯一值及其索引是使用cython 的显式循环找到的
-
单程方法
-
cached_dict()：使用 Python dict 进行缓存（O(1) 查找）
-
cached_dict_cy()：与上述相同，但使用 Cython（基本上等同于来自 here 的 vectorized_cached_impl()）
-
cached_arr_cy()：使用数组进行缓存（O(U) 查找）
纯()
def pure(x):
return 2 * x
np.vectorized()
import numpy as np
vectorized = np.vectorize(pure)
vectorized.__name__ = 'vectorized'
np_unique()
import functools
import numpy as np
def vectorize_np_unique(func):
@functools.wraps(func)
def func_vect(arr):
uniques, ix = np.unique(arr, return_inverse=True)
result = np.array([func(x) for x in uniques])
return result[ix].reshape(arr.shape)
return func_vect
np_unique = vectorize_np_unique(pure)
np_unique.__name__ = 'np_unique'
pd_unique()
import functools
import numpy as np
import pandas as pd
def vectorize_pd_unique(func):
@functools.wraps(func)
def func_vect(arr):
shape = arr.shape
arr = arr.ravel()
uniques = np.sort(pd.unique(arr))
f_range = np.array([func(x) for x in uniques])
return f_range[np.searchsorted(uniques, arr)].reshape(shape)
return func_vect
pd_unique = vectorize_pd_unique(pure)
pd_unique.__name__ = 'pd_unique'
set_unique()
import functools
def vectorize_set_unique(func):
@functools.wraps(func)
def func_vect(arr):
shape = arr.shape
arr = arr.ravel()
uniques = sorted(set(arr))
result = np.array([func(x) for x in uniques])
return result[np.searchsorted(uniques, arr)].reshape(shape)
return func_vect
set_unique = vectorize_set_unique(pure)
set_unique.__name__ = 'set_unique'
set_unique_msk()
import functools
def vectorize_set_unique_msk(func):
@functools.wraps(func)
def func_vect(arr):
result = np.empty_like(arr)
for x in set(arr.ravel()):
result[arr == x] = func(x)
return result
return func_vect
set_unique_msk = vectorize_set_unique_msk(pure)
set_unique_msk.__name__ = 'set_unique_msk'
nb_unique()
import functools
import numpy as np
import numba as nb
import flyingcircus as fc
@nb.jit(forceobj=False, nopython=True, nogil=True, parallel=True)
def numba_unique(arr, max_uniques):
ix = np.empty(arr.size, dtype=np.int64)
uniques = np.empty(max_uniques, dtype=arr.dtype)
j = 0
for i in range(arr.size):
found = False
for k in nb.prange(j):
if arr[i] == uniques[k]:
found = True
break
if not found:
uniques[j] = arr[i]
j += 1
uniques = np.sort(uniques[:j])
# : get indices
num_uniques = j
for j in nb.prange(num_uniques):
x = uniques[j]
for i in nb.prange(arr.size):
if arr[i] == x:
ix[i] = j
return uniques, ix
@fc.base.parametric
def vectorize_nb_unique(func, max_uniques=-1):
@functools.wraps(func)
def func_vect(arr):
nonlocal max_uniques
shape = arr.shape
arr = arr.ravel()
if max_uniques <= 0:
m = arr.size
elif isinstance(max_uniques, int):
m = min(max_uniques, arr.size)
elif isinstance(max_uniques, float):
m = int(arr.size * min(max_uniques, 1.0))
uniques, ix = numba_unique(arr, m)
result = np.array([func(x) for x in uniques])
return result[ix].reshape(shape)
return func_vect
nb_unique = vectorize_nb_unique()(pure)
nb_unique.__name__ = 'nb_unique'
cy_unique()
%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True
import numpy as np
import cython as cy
cimport cython as ccy
cimport numpy as cnp
ctypedef fused arr_t:
cnp.uint16_t
cnp.uint32_t
cnp.uint64_t
cnp.int16_t
cnp.int32_t
cnp.int64_t
cnp.float32_t
cnp.float64_t
cnp.complex64_t
cnp.complex128_t
def sort_numpy(arr_t[:] a):
np.asarray(a).sort()
cpdef cnp.int64_t cython_unique(
arr_t[:] arr,
arr_t[::1] uniques,
cnp.int64_t[:] ix):
cdef size_t size = arr.size
cdef arr_t x
cdef cnp.int64_t i, j, k, num_uniques
j = 0
for i in range(size):
found = False
for k in range(j):
if arr[i] == uniques[k]:
found = True
break
if not found:
uniques[j] = arr[i]
j += 1
sort_numpy(uniques[:j])
num_uniques = j
for j in range(num_uniques):
x = uniques[j]
for i in range(size):
if arr[i] == x:
ix[i] = j
return num_uniques
import functools
import numpy as np
import flyingcircus as fc
@fc.base.parametric
def vectorize_cy_unique(func, max_uniques=0):
@functools.wraps(func)
def func_vect(arr):
shape = arr.shape
arr = arr.ravel()
if max_uniques <= 0:
m = arr.size
elif isinstance(max_uniques, int):
m = min(max_uniques, arr.size)
elif isinstance(max_uniques, float):
m = int(arr.size * min(max_uniques, 1.0))
ix = np.empty(arr.size, dtype=np.int64)
uniques = np.empty(m, dtype=arr.dtype)
num_uniques = cy_uniques(arr, uniques, ix)
uniques = uniques[:num_uniques]
result = np.array([func(x) for x in uniques])
return result[ix].reshape(shape)
return func_vect
cy_unique = vectorize_cy_unique()(pure)
cy_unique.__name__ = 'cy_unique'
cached_dict()
import functools
import numpy as np
def vectorize_cached_dict(func):
@functools.wraps(func)
def func_vect(arr):
result = np.empty_like(arr.ravel())
cache = {}
for i, x in enumerate(arr.ravel()):
if x not in cache:
cache[x] = func(x)
result[i] = cache[x]
return result.reshape(arr.shape)
return func_vect
cached_dict = vectorize_cached_dict(pure)
cached_dict.__name__ = 'cached_dict'
cached_dict_cy()
%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True
import numpy as np
import cython as cy
cimport cython as ccy
cimport numpy as cnp
ctypedef fused arr_t:
cnp.uint16_t
cnp.uint32_t
cnp.uint64_t
cnp.int16_t
cnp.int32_t
cnp.int64_t
cnp.float32_t
cnp.float64_t
cnp.complex64_t
cnp.complex128_t
ctypedef fused result_t:
cnp.uint16_t
cnp.uint32_t
cnp.uint64_t
cnp.int16_t
cnp.int32_t
cnp.int64_t
cnp.float32_t
cnp.float64_t
cnp.complex64_t
cnp.complex128_t
cpdef void apply_cached_dict_cy(arr_t[:] arr, result_t[:] result, object func):
cdef size_t size = arr.size
cdef size_t i
cdef dict cache = {}
cdef arr_t x
cdef result_t y
for i in range(size):
x = arr[i]
if x not in cache:
y = func(x)
cache[x] = y
else:
y = cache[x]
result[i] = y
import functools
import flyingcircus as fc
@fc.base.parametric
def vectorize_cached_dict_cy(func, dtype=None):
@functools.wraps(func)
def func_vect(arr):
nonlocal dtype
shape = arr.shape
arr = arr.ravel()
result = np.empty_like(arr) if dtype is None else np.empty(arr.shape, dtype=dtype)
apply_cached_dict_cy(arr, result, func)
return np.reshape(result, shape)
return func_vect
cached_dict_cy = vectorize_cached_dict_cy()(pure)
cached_dict_cy.__name__ = 'cached_dict_cy'
cached_arr_cy()
%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True
import numpy as np
import cython as cy
cimport cython as ccy
cimport numpy as cnp
ctypedef fused arr_t:
cnp.uint16_t
cnp.uint32_t
cnp.uint64_t
cnp.int16_t
cnp.int32_t
cnp.int64_t
cnp.float32_t
cnp.float64_t
cnp.complex64_t
cnp.complex128_t
ctypedef fused result_t:
cnp.uint16_t
cnp.uint32_t
cnp.uint64_t
cnp.int16_t
cnp.int32_t
cnp.int64_t
cnp.float32_t
cnp.float64_t
cnp.complex64_t
cnp.complex128_t
cpdef void apply_cached_arr_cy(
arr_t[:] arr,
result_t[:] result,
object func,
arr_t[:] uniques,
result_t[:] func_uniques):
cdef size_t i
cdef size_t j
cdef size_t k
cdef size_t size = arr.size
j = 0
for i in range(size):
found = False
for k in range(j):
if arr[i] == uniques[k]:
found = True
break
if not found:
uniques[j] = arr[i]
func_uniques[j] = func(arr[i])
result[i] = func_uniques[j]
j += 1
else:
result[i] = func_uniques[k]
import functools
import numpy as np
import flyingcircus as fc
@fc.base.parametric
def vectorize_cached_arr_cy(func, dtype=None, max_uniques=None):
@functools.wraps(func)
def func_vect(arr):
nonlocal dtype, max_uniques
shape = arr.shape
arr = arr.ravel()
result = np.empty_like(arr) if dtype is None else np.empty(arr.shape, dtype=dtype)
if max_uniques is None or max_uniques <= 0:
max_uniques = arr.size
elif isinstance(max_uniques, int):
max_uniques = min(max_uniques, arr.size)
elif isinstance(max_uniques, float):
max_uniques = int(arr.size * min(max_uniques, 1.0))
uniques = np.empty(max_uniques, dtype=arr.dtype)
func_uniques = np.empty_like(arr) if dtype is None else np.empty(max_uniques, dtype=dtype)
apply_cached_arr_cy(arr, result, func, uniques, func_uniques)
return np.reshape(result, shape)
return func_vect
cached_arr_cy = vectorize_cached_arr_cy()(pure)
cached_arr_cy.__name__ = 'cached_arr_cy'
注意事项
元装饰器@parametric（灵感来自here，在FlyingCircus中作为flyingcircus.base.parametric可用）定义如下：
def parametric(decorator):
@functools.wraps(decorator)
def _decorator(*_args, **_kws):
def _wrapper(func):
return decorator(func, *_args, **_kws)
return _wrapper
return _decorator
Numba 无法比常规 Python 代码更有效地处理 单次传递 方法，因为传递任意 callable 需要启用 Python object 支持，从而排除快速 JIT 循环。
Cython 有一些限制，您需要指定预期的结果数据类型。您也可以从输入数据类型中试探性地猜测它，但这并不理想。
为了简单起见，使用静态 NumPy 数组实现了一些需要临时存储的实现。例如，可以使用 C++ 中的动态数组来改进这些实现，而不会造成太大的速度损失，但会大大改善内存占用。
基准测试
只有 10 个唯一值的慢速函数（小于 ~0.05%）
（这基本上是原始帖子的用例）。
具有 ~0.05% 唯一值的快速函数
具有 ~10% 唯一值的快速函数
具有 ~20% 唯一值的快速函数
完整的基准代码（基于this template）在here 上可用。
讨论与结论
最快的方法取决于N 和U。
对于慢速函数，所有缓存方法都比vectorized() 更快。当然，这个结果应该持保留态度，因为这里测试的 slow 函数比 fast 函数慢了约 4 个数量级，而且这种慢速分析函数并不太常见。
如果函数可以立即以矢量化形式编写，那是迄今为止最快的方法。
一般来说，只要U / N 小于或等于 20%，cached_dict_cy() 的内存效率和速度都比 vectorized() 快（即使对于快速函数也是如此）。
它的主要缺点是需要 Cython，这是一个有点复杂的依赖项，并且还需要指定结果数据类型。
只要U / N 小于或等于 10%，np_unique() 方法就比vectorized() 快（即使对于快速函数也是如此）。
pd_unique() 方法仅对非常小的 U 和慢功能具有竞争力。
对于非常小的U，散列的好处略少，cached_arr_cy() 是最快的方法。