使用 if 和 for 的嵌套列表理解

Question

我試圖使用numpy輕松輸出數組的正索引。 我能夠得到以下信息：

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]

這給了我數組a_pos_idx數組a的正索引。 但是，我想使用列表理解更簡單地（以可讀的單行方式）執行此操作。 這就是我想出的：

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

但是，這給了我一個無效的語法錯誤。 有沒有辦法在 Python 中創建這樣的嵌套 for 循環？

Answer 1

列表推導式需要將過濾器（ if a[i] > 0 ）放在最后。

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

由於您使用的是 numpy，您可以使用內置的where函數，如下所示：

np.where(a > 0)

完整示例

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}")

輸出：

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]