python_learner
python_learner

Reputation: 27

List comprehensions that get rid of the corresponding elements in parallel lists

I want to get rid of all of the zeros in an array and in another array, get rid of the elements corresponding to the indices of the ones that did have zeros in the first array.

I have:

A = [2,3,2,4,5]
P = [0.1,0.3,0.4,0.5,0]
new_P = [x for x in P if P[x]!=0]
new_A = [x for x in A if P[x]!=0]

What am I doing wrong? For new_P, for example, I get a blank array.

Upvotes: 0

Views: 80

Answers (2)

Chris
Chris

Reputation: 22973

  • Your first list comprehension will not work because your iterating by reference, but you're trying to access an element by index. Python is going to raise an IndexError. Change your comparison to test with x instead of P[x]: new_P = [x for x in P if x != 0].

  • Your second list comprehension will not work for the exact reason the first will not work. But the same fix does not apply to your second list comprehension. You need to use zip to iterate over both P and A, testing if the current element in P is a zero.

    >>> A = [2,3,2,4,5]
    >>> P = [0.1,0.3,0.4,0.5,0]
    >>> new_P = [el for el in P if el != 0] # testing using x not P[x]
    >>> new_P
    [0.1, 0.3, 0.4, 0.5]
    >>> new_A = [el[0] for el in zip(A, P) if el[1] != 0]
    >>> new_A
    [2, 3, 2, 4]
    >>> 
    

Upvotes: 1

pzp
pzp

Reputation: 6607

How about just zipping the lists together and processing their elements together?

a = [2, 3, 2, 0, 5]
p = [0.1, 0.3, 0.4, 0.5, 0.6]

a, p = zip(*[(a, p) for a, p in zip(a, p) if p != 0])

Upvotes: 0

Related Questions