如何利用topk获取分布式Dask数组的全局最小n个值？

你提出的分两次调用topk分别针对不同轴的思路是完全可行的，而且特别适合Dask的分布式计算场景——它能巧妙利用数组的分块结构，逐步缩小候选值范围，避免一次性处理全部数据带来的性能问题。

为什么两次topk的方法有效？ #

Dask的topk操作是分布式执行的：它会先在每个数据块上计算局部的topk结果，再将这些局部结果合并得到全局的轴方向topk。针对你的需求，我们可以通过两次轴方向的筛选，把全局最小的n个值锁定在一个小的候选集中：

1. 第一次调用topk(-m, axis=0)：在每一列中筛选出最小的m个值（这里m建议设为比目标n大的数，比如m = n*5，避免因局部筛选遗漏全局最小值），得到一个形状为(m, 2400)的数组。
2. 第二次调用topk(-m, axis=1)：在第一步得到的数组的每一行中，再筛选出最小的m个值，得到一个(m, m)的候选数组——这个数组里已经包含了全局最小的n个值。
3. 最后将候选数组扁平化，再调用一次topk(-n)就能得到整个数组中最小的n个值。

具体代码示例 #

import dask.array as da

# 假设你的分布式数组是dist，形状(2400,2400)，块大小(100,100)
dist = da.random.random((2400,2400), chunks=(100,100))

n = 5  # 你要找的最小的n个值
m = n * 5  # 候选集大小，可根据实际情况调整

# 第一步：按列筛选最小的m个值
step1 = dist.topk(-m, axis=0)
# 第二步：按行筛选最小的m个值
step2 = step1.topk(-m, axis=1)
# 第三步：扁平化候选集，取全局最小的n个值
global_topn = step2.flatten().topk(-n).compute()

print(global_topn)

另一种思路：直接扁平化数组 #

如果你处理的数组规模不算特别大，也可以直接将数组扁平化后调用topk，代码更简洁：

global_topn_flat = dist.flatten().topk(-n).compute()

不过这种方法的缺点是，扁平化操作可能会触发更多跨节点的数据传输，对于超大型分布式数组，性能不如两次topk的筛选方法。

补充说明 #

你最初的代码得到的是一个5x5的二维数组，这只是经过两次轴筛选后的候选集，里面包含了全局最小的5个值，但需要再扁平化后取topk，才能得到排序好的全局最小n个值哦。

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