如何替换 Counter 以仅使用 numpy 代码

Question

我有这个代码：

from collections import Counter
import numpy as np

def make_data(N):
    np.random.seed(40)
    g = np.random.randint(-3, 4, (N, N))
    return g


N = 100
g = make_data(N)
n = g.shape[0]

sum_dist = Counter()
for i in range(n):
    for j in range(n):
        dist = i**2 + j**2
        sum_dist[dist] += g[i, j]

sorted_dists = sorted(sum_dist.keys())
for i in range(1, len(sorted_dists)):
    sum_dist[sorted_dists[i]] += sum_dist[sorted_dists[i-1]]

# print(sum_dist)
print(max(sum_dist, key=sum_dist.get))

输出为 7921。

我想将其转换为仅 numpy 代码并摆脱 Counter。我怎样才能做到这一点？

Answer 1

sum_dist既然你知道它的最大索引，你能直接把它变成一个数组吗？

sum_dist = np.zeros(2 * N * N, dtype=int)
for i in range(n):
    for j in range(n):
        dist = i**2 + j**2
        sum_dist[dist] += g[i, j]

print(np.argmax(np.cumsum(sum_dist)))