How can I use numpy.sum to output sum as 2d array without a for loop?

Question

I would like to get the sum of each teams score in each round, currently I am able to get the sum of each teams scores from round 1.

Current Input

scores = np.array([
    [1, 2, 2, 3, 5, 8, 12], #Round 1
    [11, 3, 9, 2, 3, 5, 10]] # Round 2
)

teams = np.array([[0, 1, 2], [1, 2, 3], [6, 5, 4]])

np.sum(np.take(scores, teams), axis=1)

This outputs the correct sum for each team in the round 1

array([ 5,  7, 25])

Is there a way to make it output the sum for each team in each round without using a for loop?

Desired Output

array([[ 5,  7, 25], [23, 14, 18]])

the rows in teams are the teams so each team has 3 players. The rows in scores are the scores of each player by index so player 0's score in Round 1 was 1 scores[0][0] — radio23
– radio23, Commented Feb 20, 2022 at 19:55

user17242583 · Accepted Answer · 2022-02-20 19:36:58Z

3

You can add axis=1 to np.take, and then sum on axis 2:

>>> np.take(scores, teams, axis=1).sum(axis=2)
array([[ 5,  7, 25],
       [23, 14, 18]])

answered Feb 20, 2022 at 19:36

user17242583

Sign up to request clarification or add additional context in comments.

1 Comment

radio23 Over a year ago

This works, I was under the impression that removing a for loop would speed up my code but it seems just a for loop and indexing into scores is faster when testing with 100,000 teams and 10,000 scores scores = np.random.rand(10_000,100) teams = np.array([random.sample(list(range(100)), 8) for _ in range(100_000)])

hpaulj · Accepted Answer · 2022-02-20 19:36:53Z

1

You didn't specify an axis, so take, as documented, works with the flattened array:

In [216]: np.take(scores.ravel(),teams)
Out[216]: 
array([[ 1,  2,  2],
       [ 2,  2,  3],
       [12,  8,  5]])

Using teams to index the columns:

In [220]: scores[:,teams]
Out[220]: 
array([[[ 1,  2,  2],
        [ 2,  2,  3],
        [12,  8,  5]],

       [[11,  3,  9],
        [ 3,  9,  2],
        [10,  5,  3]]])

and summing:

In [221]: scores[:,teams].sum(axis=2)
Out[221]: 
array([[ 5,  7, 25],
       [23, 14, 18]])

answered Feb 20, 2022 at 19:36

hpaulj

233k14 gold badges260 silver badges392 bronze badges

1 Comment

radio23 Over a year ago

This solution works, but when run with larger arrays it becomes very slow, Im guessing due to memory issues?

Collectives™ on Stack Overflow

How can I use numpy.sum to output sum as 2d array without a for loop?

Current Input

Desired Output

2 Answers 2

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

Current Input

Desired Output

2 Answers 2

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related