CODEWITHSUNDEEP
pythonarraysnumpy
Shan
Shan

Reputation: 19263

Setting value of numpy array based on multiple criteria

I am trying to set the values in a numpy array to zero if it is equivalent to any number in a list.

Lets consider the following array

  a = numpy.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])

I want to set multiple elements of a which are in the list [1, 2, 8] to 0.

The result should be 

   [[0, 0, 3],
    [4, 0, 6],
    [7, 0, 9]]

For a single element it's simple

   a[a == 1] = 0

The above only works for a single integer. How it could work for a list?

Upvotes: 4

Views: 2166

Answers (4)

Spencer Hill
Spencer Hill

Reputation: 899

Combining my comment to the original question regarding np.where and the excellent answer by @Jaime above using np.in1d:

import numpy as np
a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
a = np.where(np.in1d(a, [1,2,8]).reshape(a.shape), 0, a)

EDIT Looks like Jaime's solution is slightly faster:

In [3]: %timeit a[np.in1d(a, [1, 2, 8]).reshape(a.shape)] = 0
10000 loops, best of 3: 45.8 µs per loop
In [4]: %timeit np.where(np.in1d(a, [1,2,8]).reshape(a.shape), 0, a)
10000 loops, best of 3: 66.7 µs per loop

Upvotes: 2

user2555451
user2555451

Reputation:

You can use numpy.vectorize:

>>> import numpy
>>> a = numpy.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
>>> lst = [1, 2, 8]
>>> a[numpy.vectorize(lambda x: x in lst)(a)] = 0
>>> a
array([[0, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])
>>>

Upvotes: 1

Jaime
Jaime

Reputation: 67457

Using np.in1d you could do the following:

>>> a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
>>> np.in1d(a, [1, 2, 8])
array([ True,  True, False, False,  True, False, False,  True, False], dtype=bool)
>>> a[np.in1d(a, [1, 2, 8]).reshape(a.shape)] = 0
>>> a
array([[0, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])

Upvotes: 5

qwwqwwq
qwwqwwq

Reputation: 7329

You can make your own vectorized functions using numpy.vectorize:

import numpy as np
a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])

def is_target_number(x):
    if x in set([1,2,8]):
        return True
    else:
        return False

f = np.vectorize(is_target_number)

a[f(a)] = 0

A lot of operators like the equality operator are already vectorized by default, numpy.vectorize allows you to use more complicated logic with the same benefit of succinctness. And if you're into code golf you can do something like this:

a[np.vectorize(lambda x: (x in set([1,2,8])))(a)] = 0

Upvotes: 1

Related Questions