CODEWITHSUNDEEP
pythonnumpylist-comprehension
Skeleton Bow
Skeleton Bow

Reputation: 529

Nested list comprehension using if and for

I was trying to easily output the positive indices of an array using numpy. I was able to get the following:

import numpy as np

np.random.seed(1)
a = np.random.sample(100) - 0.5
a_pos_idx = []
for i in range(len(a)):
  if a[i] > 0:
    a_pos_idx += [i]

This gives me the positive indices in the array a inside the array a_pos_idx. However, I wanted to do this more simply (in a readable one-liner) using list comprehension. This is what I came up with:

a_pos_idx = [i if a[i] > 0 for i in range(len(a))]

However, this gives me an invalid syntax error. Is there any way to create such nested for loops in Python?

Upvotes: 0

Views: 63

Answers (1)

Christopher Peisert
Christopher Peisert

Reputation: 24134

The list comprehension needs to place the filter (if a[i] > 0) at the end.

[i for i in range(len(a)) if a[i] > 0]

Since you are using numpy, you can use the built-in where function as follows:

np.where(a > 0)

Full example

import numpy as np

np.random.seed(1)
a = np.random.sample(100) - 0.5
a_pos_idx = []
for i in range(len(a)):
  if a[i] > 0:
    a_pos_idx += [i]

print(f"For loop:\n{a_pos_idx}")

a_pos_idx_2 = [i for i in range(len(a)) if a[i] > 0]
print(f"\nList comphrehension:\n{a_pos_idx_2}")

a_pos_idx_3, *_ = np.where(a > 0)
print(f"\nnumpy filter:\n{a_pos_idx_3}")

Output:

For loop:
[1, 9, 11, 13, 15, 17, 20, 21, 23, 24, 25, 29, 32, 33, 34, 36, 37, 39, 40, 41, 43, 46, 51, 56, 58, 59, 62, 65, 66, 67, 68, 69, 70, 73, 76, 78, 79, 80, 81, 82, 85, 87, 88, 89, 91, 93, 96, 97, 99]

List comphrehension:
[1, 9, 11, 13, 15, 17, 20, 21, 23, 24, 25, 29, 32, 33, 34, 36, 37, 39, 40, 41, 43, 46, 51, 56, 58, 59, 62, 65, 66, 67, 68, 69, 70, 73, 76, 78, 79, 80, 81, 82, 85, 87, 88, 89, 91, 93, 96, 97, 99]

numpy filter:
[ 1  9 11 13 15 17 20 21 23 24 25 29 32 33 34 36 37 39 40 41 43 46 51 56
 58 59 62 65 66 67 68 69 70 73 76 78 79 80 81 82 85 87 88 89 91 93 96 97
 99]

Upvotes: 2

Related Questions