选择基于where语句的行

Posted 2023-01-25

如何选择在其中包含“链接”一词的值，并将它们放在category1中，并在其中添加“爆米花”以使其成为类别2而将所有其他内容放入类别3中？

这是一个示例，但我的实际数据集有数百行

data = 'model': [['Lisa', 'link'], ['Lisa 2', 'popcorn'], ['telephone', 'rabbit']],
        'launched': [1983, 1984, 1991]

df = pd.DataFrame(data, columns = ['model', 'launched'])

期望

 Model                 launched         category
 ['Lisa', 'link']        1983             1
 ['Lisa 2', 'popcorn']   1984             2
 ['telephone', 'rabbit'] 1991             3

答案

你可以使用apply函数：

创建一个def：

def get_categories(row):
    if 'link' in row.model:
        return 1
    elif 'popcorn' in row.model:
        return 2
    else:
        return 3

然后像这样称呼它：

df['category'] = df.apply(get_categories, axis=1)
df

输出：

    model           launched    category
0   [Lisa, link]        1983    1
1   [Lisa 2, popcorn]   1984    2
2   [telephone, rabbit] 1991    3

编辑：

基于@gred_data注释，您实际上可以在一行中执行此操作以提高性能：

df['category'] = df.model.apply(lambda x: 1 if 'link' in x else 2 if 'popcorn' in x else 3)
df

获得相同的结果。

另一答案

您可以使用np.select将category设置为1或2，具体取决于'link'或'popcorn'是否包含在给定列表中。将default设置为3，以确定它们都不包含在内：

import numpy as np
c1 = ['link' in i for i in df.model]
c2 = ['popcorn' in i for i in df.model]
df['category'] = np.select([c1,c2], [1,2], 3)

              model       launched  category
0         [Lisa, link]      1983         1
1    [Lisa 2, popcorn]      1984         2
2  [telephone, rabbit]      1991         3