如何移除Counter依赖，仅用Numpy实现指定Python代码功能

现有如下Python代码（使用了Counter和Numpy库）：

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。现需移除Counter依赖，仅使用Numpy库实现相同功能，实现方式如下：

纯Numpy实现方案 #

核心思路是用Numpy的向量化操作替代循环和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)

# 利用广播机制生成所有(i,j)对应的dist值：i²+j²
i = np.arange(N)[:, np.newaxis]  # 转为列向量，形状(N,1)
j = np.arange(N)[np.newaxis, :]  # 转为行向量，形状(1,N)
dist = i**2 + j**2

# 扁平化二维数组，方便后续分组操作
dist_flat = dist.flatten()
g_flat = g.flatten()

# 获取所有唯一的dist值并按升序排序
unique_dists = np.unique(dist_flat)

# 对每个唯一dist值，计算对应g元素的总和
sum_vals = np.array([g_flat[dist_flat == d].sum() for d in unique_dists])

# 计算前缀和，对应原代码中累加前序sum_dist的逻辑
prefix_sum = np.cumsum(sum_vals)

# 找到前缀和最大的那个dist值
max_dist = unique_dists[np.argmax(prefix_sum)]

print(max_dist)

关键步骤说明： #

• 广播生成dist矩阵：避免嵌套循环，一次性生成所有坐标对应的距离值，效率远高于循环遍历。
• 分组求和替代Counter：通过np.unique获取唯一距离值，再针对每个值筛选对应的g元素求和，实现Counter的累加功能。
• 前缀和简化循环：用np.cumsum直接计算前缀和，替代原代码中手动循环累加前序值的逻辑，代码更简洁高效。

运行这段代码，输出结果同样为7921，与原代码功能完全一致。

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