Home > Software design >  Count number of elements in each row equal to the first element in each row
Count number of elements in each row equal to the first element in each row

Time:03-05

a = np.array((['1', '1', '1', '2'],
              ['3', '3', '3', '3'], 
              ['2', '1', '1', '2'], 
              ['1', '3', '1', '2']))

I am looking for a way (rather than iterating over the array) to count the number of elements in each row that are the equal to the first element of the row. i.e. in the first row, the first element is '1' and there are 2 other occurrences of this value. The final result should be:

result = [2, 3, 1, 1]

My array can be quite large so iterating over each row will be slow.

CodePudding user response:

You can do it using List Comprehension:

result = [list(a[i,:]).count(a[i,0])-1 for i in range(a.shape[0])]
print(result)

This gives me:

[2, 3, 1, 1]

CodePudding user response:

to do this, you can take the first column as a view using slicing notation, then take the remaining sections as a view, test for equality (numpy will "broadcast" arrays that are of different dimensions if they are compatible), then perform a row-wise count.

equalities = a[:,:1] == a[:,1:]
result = equalities.sum(axis=1)

this gives a numpy 1d array containing the result:

array([2, 3, 1, 1])

CodePudding user response:

Compare to the slice a[..., 0]. Append a new axis to make it the correct dimensions again.

result = (a == a[..., 0][..., np.newaxis]).sum(axis=-1) - 1

The ... ensure that this works for arrays of any dimension, whereas using : only works for 2D.

  • Related