[FIXED] How to replace Counter to use numpy code only

Issue

I have this code:

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))

The output is 7921.

I want to convert it into numpy only code and get rid of Counter. How can I do that?

Solution

Can you just make sum_dist into an array, since you know its maximum index?

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)))

Answered By - Frank Yellin