Michael
Michael

Reputation: 367

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

I have a 5x5 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?

Upvotes: 0

Views: 127

Answers (1)

user2999345
user2999345

Reputation: 4195

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()

Upvotes: 3

Related Questions