CODEWITHSUNDEEP
python-3.xnumpy
Bruce
Bruce

Reputation: 585

how to assign labels of one numpy array to another numpy array and group accordingly?

I have the following labels 

>>> lab
array([2, 2, 2, 2, 2, 3, 3, 0, 0, 0, 0, 1])

I want to assign this label to another numpy array i.e

>>> arr
array([[81,  1,  3, 87],  # 2
       [ 2,  0,  1,  0],  # 2
       [13,  6,  0,  0],  # 2
       [14,  0,  1, 30],  # 2
       [ 0,  0,  0,  0],  # 2
       [ 0,  0,  0,  0],  # 3
       [ 0,  0,  0,  0],  # 3
       [ 0,  0,  0,  0],  # 0
       [ 0,  0,  0,  0],  # 0
       [ 0,  0,  0,  0],  # 0 
       [ 0,  0,  0,  0],  # 0
       [13,  2,  0, 11]]) # 1

and add the elements of 0th group, 1st group, 2nd group, 3rd group?

Upvotes: 1

Views: 1186

Answers (1)

behzad.nouri
behzad.nouri

Reputation: 77991

If the labels of equal values are contiguous, as in your example, then you may use np.add.reduceat:

>>> lab
array([2, 2, 2, 2, 2, 3, 3, 0, 0, 0, 0, 1])
>>> idx = np.r_[0, 1 + np.where(lab[1:] != lab[:-1])[0]]
>>> np.add.reduceat(arr, idx)
array([[110,   7,   5, 117],   # 2
       [  0,   0,   0,   0],   # 3
       [  0,   0,   0,   0],   # 0
       [ 13,   2,   0,  11]])  # 1

if they are not contiguous, then use np.argsort to align the array and labels such that labels of the same values are next to each other:

>>> i = np.argsort(lab)
>>> lab, arr = lab[i], arr[i, :] # aligns array and labels such that labels
>>> lab                          # are sorted and equal labels are contiguous
array([0, 0, 0, 0, 1, 2, 2, 2, 2, 2, 3, 3])
>>> idx = np.r_[0, 1 + np.where(lab[1:] != lab[:-1])[0]]
>>> np.add.reduceat(arr, idx)
array([[  0,   0,   0,   0],  # 0
       [ 13,   2,   0,  11],  # 1
       [110,   7,   5, 117],  # 2
       [  0,   0,   0,   0]]) # 3

or alternatively use groupby from pandas library:

>>> pd.DataFrame(arr).groupby(lab).sum().values
array([[  0,   0,   0,   0],
       [ 13,   2,   0,  11],
       [110,   7,   5, 117],
       [  0,   0,   0,   0]])

Upvotes: 1

Related Questions