Delete columns from numpy array depending on a condition on a single cell
With numpy.delete
:
a = np.array([[0, 1, 2, 0, 4, 5, 6, 7, 8, 9], [0, 11, 0, 13, 0, 15, 0, 17, 18, 0]])
indices = [i for (i,v) in enumerate(a[1]) if v==0]
# [0, 2, 4, 6, 9]
a = np.delete(a, indices, 1)
# array([[ 1, 0, 5, 7, 8], [11, 13, 15, 17, 18]])
Simple (fully numpy) solution:
import numpy as np
t = np.array([[ 0, 1, 2, 0, 4, 5, 6, 7, 8, 9], [ 0, 11, 0, 13, 0, 15, 0, 17, 18, 0]])
indices_to_keep = t[1].nonzero()[0]
print t[:,indices_to_keep]
# [[ 1 0 5 7 8]
# [11 13 15 17 18]]
Using np.where
:
>>> t.T[np.where(t[1])].T
array([[ 1, 0, 5, 7, 8],
[11, 13, 15, 17, 18]])
Similar to Juh_, but more expressive, and avoiding some minor unnecessary performance overhead. A grand total of 12 highly pythonic, explicit and unambigious characters. This is really numpy 101; if you are still trying to wrap your head around this, you would do yourself a favor by reading a numpy primer.
import numpy as np
a = np.array([[ 0, 1, 2, 0, 4, 5, 6, 7, 8, 9],
[ 0, 11, 0, 13, 0, 15, 0, 17, 18, 0]])
print a[:,a[1]!=0]