Numpy 如何将dask数组中小于前k的所有项归零

我想将dask.array的所有元素归零，除了最上面的几个元素。我该怎么做

例子 假设我有一个小dask数组，如下所示：

import numpy as np
import dask.array as da
x = np.array([0, 4, 2, 3, 1])
x = da.from_array(x, chunks=(2,))

>>> result.compute()
array([0, 4, 0, 3, 0])

除了两个最大的元素外，我如何将所有元素归零？我想要像下面这样的东西：

import numpy as np
import dask.array as da
x = np.array([0, 4, 2, 3, 1])
x = da.from_array(x, chunks=(2,))

>>> result.compute()
array([0, 4, 0, 3, 0])

---

您可以通过函数和就地设置项的组合来实现这一点

top = x.topk(2)
x[x < top[-1]] = 0

>>> x.compute()
array([0, 4, 0, 3, 0])

只有当您试图在一台机器上通过一个大型数据集进行流式传输时，这一点才有意义。如果您在分布式系统上，这一点不会对您产生太大影响