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

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

CodePudding user response：

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

CodePudding user response：

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