CODEWITHSUNDEEP
pythonarraysnumpy
peremeykin
peremeykin

Reputation: 549

Obtain references to numpy array elements by condition

So, I have a target numpy array, say 

a = np.array([3, 7, 9, 3, 5])

and I have a condition array, e.g.

c = np.array([False, False, True, True, False])

and I'd like to obtain array b, specified by c 

>>> b
array([9,3])

such that if I change it, a also changes. For example:

>>> b[0]=2
>>> b
array([2,3])
>>> a
array([3, 7, 2, 3, 5])

I've tried something like that:

>>> b = np.compress(c,a)
>>> b
array([9, 3])
>>> b[0]=2
>>> b
array([2, 3])
>>> a
array([3, 7, 9, 3, 5])

But it does not work, a is still same, np.compress() returns a copy. Also I had no success with this code:

b=a[np.where(c)]

Is there any analog of np.compress() that returns references, not copy? Thank you!

Upvotes: 2

Views: 120

Answers (1)

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 96172

One thing that could probably fit your needs is numpy "masking":

>>> import numpy as np
>>> a = np.array([3,7, 9, 3, 5])
>>> c = np.array([False, False, True, True, False])
>>> b = np.ma.array(a,mask=~c)

Note I had to take the elementwise negation of c: ~c, since True will be interpreted as "this element should be masked" wheras you want the opposite - this element should be unmasked. Notice now that b acts how you want when you do operations on it:

>>> b
masked_array(data = [-- -- 9 3 --],
             mask = [ True  True False False  True],
       fill_value = 999999)
>>> b.sum()
12
>>> (2*b)
masked_array(data = [-- -- 18 6 --],
             mask = [ True  True False False  True],
       fill_value = 999999)

>>> (2*b).sum()

And furthermore, it is working on a view of a, so changing a changes b:

>>> a
array([3, 7, 9, 3, 5])
>>> a[2] = 42
>>> a
array([ 3,  7, 42,  3,  5])
>>> b
masked_array(data = [-- -- 42 3 --],
             mask = [ True  True False False  True],
       fill_value = 999999)

>>> b.sum()
45

The reverse also work, albeit, the index doesn't change:

>>> b[2] = 88
>>> a
array([ 3,  7, 88,  3,  5])
>>> b
masked_array(data = [-- -- 88 3 --],
             mask = [ True  True False False  True],
       fill_value = 999999)

>>>

Unfortunately, python is relatively high-level, so working with raw references is not straightforward.

Upvotes: 2

Related Questions