Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Neat way to conditionally select elements from two lists?

Tags:

python

list

I have two lists:

a = [-1, 2, 3]
b = ['a', 'b', 'c']

The number of elements in a and b are always the same.

I want to keep all the elements in a that are positive, and the corresponding elements of b.

One straightforward way would be

anew = []
bnew = []
for i in xrange(len(a)):
    if a[i] > 0:
         anew.append(a[i])
         bnew.append(b[i])

or something like

tmp = zip(*[(a[i], b[i]) for i in xrange(len(a)) if a[i]>0])
a = tmp[0] # The original a and b need not to be kept.
b = tmp[1]

I wonder if there is any shorter and cleaner (without any temporary variables) way for doing this. Some kind of in-place deletion would be perfect, but if I use del a[i], b[i] in a loop, the index will be wrong after deletion of one element.

EDIT: What if the selection involves two lists, for example if the condition becomes

if a[i] > 0 and b[i] > 0?

Maybe in this case a temporary variable has to be used?

like image 311
FJDU Avatar asked Mar 06 '14 05:03

FJDU


1 Answers

If you use numpy, there is a quite neat way:

In [117]: a = np.array([-1, 2, 3])
     ...: b = np.array(['a', 'b', 'c'])

In [119]: a[a>0]
Out[119]: array([2, 3])

In [120]: b[a>0]
Out[120]: 
array(['b', 'c'], 
      dtype='|S1')

If without numpy, you can use list comprehension with conditional statement:

In [121]: a = [-1, 2, 3]
     ...: b = ['a', 'b', 'c']

In [122]: print [i for i in a if i>0]
[2, 3]

In [124]: print [v for i, v in enumerate(b) if a[i]>0]
['b', 'c']
like image 122
zhangxaochen Avatar answered Sep 25 '22 05:09

zhangxaochen