Skip to content
Advertisement

How to get all values in an array whose added index size is greater than value?

I have a 5×5 ndarray and want to sum up all values whose added index size is greater than a given value.

For example, I have the following array

x = np.random.random([5, 5])

and want to sum all values whose row and column index is, combined, larger than 6. If I do this manually, I would calculate

idx_gr_8 = x[4, 3] + x[4, 4] + x[3, 4]

because 4 + 3, 4 + 4 and 3 + 4 are the only indices larger than 6. However this is cumbersome for larger array. Is there a numpy command or more efficient way to do this?

Advertisement

Answer

you can use meshgrid do get row and col indices:

a = np.random.rand(5, 5)
min_ind = 6

row_i, col_i = np.meshgrid(range(a.shape[0]), range(a.shape[1]), indexing='ij')
valid_mask = (row_i + col_i) > min_ind

res = a[valid_mask].sum()
User contributions licensed under: CC BY-SA
8 People found this is helpful
Advertisement