如何根据数据框列中的列表值创建多个标志列？答案

【问题标题】：How to create multiple flag columns based on list values found in the dataframe column?如何根据数据框列中的列表值创建多个标志列？
【发布时间】：2020-03-31 23:21:59
【问题描述】：

表格如下所示：

     ID  |CITY
    ----------------------------------
    1  |London|Paris|Tokyo
    2  |Tokyo|Barcelona|Mumbai|London
    3  |Vienna|Paris|Seattle

city 列包含大约 1000 多个值，它们是 |分隔

我想创建一个标志列来指示一个人是否只访问了感兴趣的城市。

    city_of_interest=['Paris','Seattle','Tokyo']

列表中有 20 个这样的值。

输出应该是这样的：

     ID      |Paris   | Seattle | Tokyo    
     -------------------------------------------
     1       |1       |0        |1      
     2       |0       |0        |1       
     3       |1       |1        |0

解决方案可以在 pandas 或 pyspark 中。

【问题讨论】：

你能把你的尝试也包括进来吗？

标签： pandas dataframe hive pyspark data-manipulation

【解决方案1】：

对于 pyspark，使用 split + array_contains:

from pyspark.sql.functions import split, array_contains

df.withColumn('cities', split('CITY', '\|')) \
  .select('ID', *[ array_contains('cities', c).astype('int').alias(c) for c in city_of_interest ]) 
  .show()
+---+-----+-------+-----+
| ID|Paris|Seattle|Tokyo|
+---+-----+-------+-----+
|  1|    1|      0|    1|
|  2|    0|      0|    1|
|  3|    1|      1|    0|
+---+-----+-------+-----+

对于 Pandas，请使用 Series.str.get_dummies:

df[city_of_interest] = df.CITY.str.get_dummies()[city_of_interest]
df = df.drop('CITY', axis=1)

【讨论】：

【解决方案2】：

熊猫解决方案

第一次转换到列表以使用DataFrame.explode：

new_df=df.copy()
new_df['CITY']=new_df['CITY'].str.lstrip('|').str.split('|')
#print(new_df)

#   ID                                CITY
#0   1              [London, Paris, Tokyo]
#1   2  [Tokyo, Barcelona, Mumbai, London]
#2   3            [Vienna, Paris, Seattle]

那么我们可以使用：

方法一：DataFrame.pivot_table

new_df=( new_df.explode('CITY')
               .pivot_table(columns='CITY',index='ID',aggfunc='size',fill_value=0)
               [city_of_interest]
               .reset_index()
               .rename_axis(columns=None)
                )
print(new_df)

方法二： DataFrame.groupby + DataFrame.unstack

new_df=( new_df.explode('CITY')
               .groupby(['ID'])
               .CITY
               .value_counts()
               .unstack('CITY',fill_value=0)[city_of_interest]
               .reset_index()
               .rename_axis(columns=None)

                )
print(new_df)

输出 new_df：

   ID  Paris  Seattle  Tokyo
0   1      1        0      1
1   2      0        0      1
2   3      1        1      0

【讨论】：

【解决方案3】：

使用 UDF 检查感兴趣的城市值是否在分隔列中。

from pyspark.sql.functions import udf

#Input list
city_of_interest=['Paris','Seattle','Tokyo']

#UDF definition
def city_present(city_name,city_list):
    return len(set([city_name]) & set(city_list.split('|')))

city_present_udf = udf(city_present,IntegerType())

#Converting cities list to a column of array type for adding columns to the dataframe
city_array = array(*[lit(city) for city in city_of_interest])
l = len(city_of_interest)
col_names = df.columns + [city for city in city_of_interest]
result = df.select(df.columns + [city_present_udf(city_array[i],df.city) for i in range(l)])
result = result.toDF(*col_names)
result.show()

【讨论】：