CODEWITHSUNDEEP
pythonarraysnumpy
taurus
taurus

Reputation: 490

Sum Rows of 2D np array with List of Indices

I have a 2d numpy array and a list of numbers. If the list is [1, 3, 1, 8] where the list sums to the number of rows, I want to output an array with the first row unchanged, the next three rows summed, the fifth row unchanged, and the remaining eight rows summed.

As an example:

A = [[0,0], [1,2], [3,4]] and l = [1, 2] would output [[0,0], [4,6]

I looked through np.sum and other functions but could not find not this functionality. Thank you.

Upvotes: 1

Views: 925

Answers (3)

taurus
taurus

Reputation: 490

I ended up coming up with my own solution when I realized I could sort my list without affecting my desired output. I used np.unique to determine the first indices of each element in the sorted list and then summed the rows between those indices. See below.

elements, indices = np.unique(data, return_counts=True)
row_summing = np.append([0], np.cumsum(indices))[:-1] #[0, index1, index2,...]

output = np.add.reduceat(matrix, row_summing, axis=0)

Upvotes: 1

hilberts_drinking_problem
hilberts_drinking_problem

Reputation: 11602

If the number of elements in l is relatively large large, you might get better performance by using groupby from pandas, e.g.

import pandas as pd

labels = np.repeat(np.arange(1, len(l) + 1), l)
# [1, 2, 2]
df = pd.DataFrame(A)
df['label'] = labels
result = df.groupby('label').sum().values

Upvotes: 1

Zabir Al Nazi Nabil
Zabir Al Nazi Nabil

Reputation: 11198

You can just iterate over the indices of l and based on the position either take that row or sum over a range of rows.

import numpy as np

A = [[0,0], [1,2], [3,4]]
l = [1, 2]

ans = []
for i in range(len(l)):
  if i%2 == 0:
    ans.append(A[ l[i] ])
  else:
    ans.append( np.sum( A[ l[i-1]:l[i-1] + l[i] ], axis=0 ) )

ans = np.array(ans)
print(ans)


[[1 2]
 [4 6]]

N.B:

If the list is [1, 3, 1, 8] where the list sums to the number of rows, I want to output an array with the first row unchanged, the next three rows summed, the fifth row unchanged, and the remaining eight rows summed.

  • I think you meant [1, 3, 5, 8]

Upvotes: 1

Related Questions