+3 votes
in Programming Languages by (17.9k points)
I want to find the indices of all elements of a 2D array that are greater than the average of the column values.

E.g. in the following matrix, it should check if 4 > average(4,1,5). If true, it should return the index of 4  "(0,0)".

     ([[4, 1, 2, 2],

       [1, 3, 9, 3],

       [5, 7, 5, 1]])

How can I do it in Pythonic way?

1 Answer

0 votes
by (31.6k points)

You can try the following code. It returns indices of all matrix elements that are greater than the column average.

>>> import numpy as np
>>> X = np.array([[4, 1, 2, 2],[1, 3, 9, 3], [5, 7, 5, 1]])
>>> X
array([[4, 1, 2, 2],
       [1, 3, 9, 3],
       [5, 7, 5, 1]])
>>> Xv=X>X.mean(axis=0)
>>> Xv
array([[ True, False, False, False],
       [False, False,  True,  True],
       [ True,  True, False, False]])
>>> idx=np.where(Xv==True)
>>> idx
(array([0, 1, 1, 2, 2]), array([0, 2, 3, 0, 1]))
>>> list(zip(idx[0], idx[1]))
[(0, 0), (1, 2), (1, 3), (2, 0), (2, 1)]

...