python - 如何在列表理解中嵌入条件？

Question

假设这个数据集:

df = pd.DataFrame({
    'name': ['John','William', 'Nancy', 'Susan', 'Robert', 'Lucy', 'Blake', 'Sally', 'Bruce', 'Mike'],
    'injury': ['right hand broken', 'lacerated left foot', 'foot broken', 'right foot fractured', '', 'sprained finger', 'chest pain', 'swelling in arm', 'laceration to arms, hands, and foot', np.NaN]
    })

    name      injury
0   John      right hand broken
1   William   lacerated left foot
2   Nancy     foot broken
3   Susan     right foot fractured
4   Robert  
5   Lucy      sprained finger
6   Blake     chest pain
7   Sally     swelling in arm
8   Bruce     lacerations to arm, hands, and foot
9   Mike      NaN
10  Jeff      swollen cheek

我只减少选定 body 部位的伤害:

selected_words = ["hand", "foot", "finger", "chest", "arms", "arm", "hands"]

df["injury"] = (
    df["injury"]
    .str.replace(",", "")
    .str.split(" ", expand=False)
    .apply(lambda x: ", ".join(set([i for i in x if i in selected_words])))
)

但是，这会向索引 9 处的 NaN 值抛出一个错误:

TypeError: 'float' object is not iterable

我将如何修改列表理解，以便:

1. 它检查任何 NaN 值

2. 如果遇到空白行或 selected_body_parts 列表中不包含正文部分(例如索引 10)，则输出 NaN

期望的输出是:

name        injury
0   John    hand
1   William foot
2   Nancy   foot
3   Susan   foot
4   Robert  NaN
5   Lucy    finger
6   Blake   chest
7   Sally   arm
8   Bruce   hand, foot, arm
9   Mike    NaN
10  Jeff    NaN

我尝试了以下方法:

.apply(lambda x: ", ".join(set([i for i in x if i in selected_words and i is not np.nan else np.nan])))

但是，语法不正确。

如有任何帮助，我们将不胜感激。谢谢!

Answer 1

您的问题不是 i 是 np.nan，而是 x 是，您无法通过理解迭代 np.nan。我想你可能想把你的 lambda 变成一个命名函数并像这样传递它:

def get_set_of_body_parts(words):
   if words is np.nan:
      return np.nan
   else:
      return ", ".join(set([i for i in x if i in selected_words]))

df = pd.DataFrame({
    'name': ['John','William', 'Nancy', 'Susan', 'Robert', 'Lucy', 'Blake', 'Sally', 'Bruce', 'Mike'],
    'injury': ['right hand broken', 'lacerated left foot', 'foot broken', 'right foot fractured', '', 'sprained finger', 'chest pain', 'swelling in arm', 'laceration to arms, hands, and foot', np.NaN]
    })

selected_words = ["hand", "foot", "finger", "chest", "arms", "arm", "hands"]

df["injury"] = (
   df["injury"]
   .str.replace(",", "")
   .str.split(" ", expand=False)
   .apply(get_set_of_body_parts)
)

但如果你真的想要，你可以像这样做一个 lambda:

.apply(lambda x: np.nan if x is np.nan else ", ".join(set([i for i in x if i in selected_words])))