如何检查列表元素是否与df中的某些标题相同？

Question

我想将列表与数据框的标题进行比较，只要所有列表元素与列标题匹配（但列表不需要包含所有列标题），代码就会继续，并删除标题与列表元素不匹配的任何列。但是，如果某些列表元素与任何列标题都不匹配，则需要用户再次输入新列表。下面是一个例子：

import pandas as pd

employees = [('Stuti', 28, 'Varanasi', 20000),
            ('Saumya', 32, 'Delhi', 25000),
            ('Aaditya', 25, 'Mumbai', 40000),
            ('Saumya', 32, 'Delhi', 35000),
            ('Saumya', 32, 'Delhi', 30000),
            ('Saumya', 32, 'Mumbai', 20000),
            ('Aaditya', 40, 'Dehradun', 24000),
            ('Seema', 32, 'Delhi', 70000)
            ]

df = pd.DataFrame(employees, columns=['Name', 'Age','City', 'Salary'], )
df = df.set_index('Name')
print(df)

         Age      City  Salary
Name
Stuti     28  Varanasi   20000
Saumya    32     Delhi   25000
Aaditya   25    Mumbai   40000
Saumya    32     Delhi   35000
Saumya    32     Delhi   30000
Saumya    32    Mumbai   20000
Aaditya   40  Dehradun   24000
Seema     32     Delhi   70000

如果假设用户输入如下列表：

user_list = ['Age', 'City']

然后 Salary 列将被删除：

         Age      City 
Name
Stuti     28  Varanasi   
Saumya    32     Delhi   
Aaditya   25    Mumbai   
Saumya    32     Delhi   
Saumya    32     Delhi   
Saumya    32    Mumbai   
Aaditya   40  Dehradun   
Seema     32     Delhi

但是如果输入的列表包含与标题不匹配的元素，代码会要求用户输入一个新的列表，如下所示：

>>> Please Enter columns you wish to keep: 
user_list = ['Age', 'Gender']
>>>'Gender' is not a column in df, please enter a new list!

本质上，列表始终是列标题列表的子集。 我想要这样的东西：

while True:
    user_list = [column for column in input('Please Enter columns you wish to keep: ').split()]
    if user_list in list(df.columns):
        # some code to drop the column that is not mentioned
        df.columns = user_list
        break
    else:
        print('try again')

Answer 1

import pandas as pd

employees = [('Stuti', 28, 'Varanasi', 20000),
            ('Saumya', 32, 'Delhi', 25000),
            ('Aaditya', 25, 'Mumbai', 40000),
            ('Saumya', 32, 'Delhi', 35000),
            ('Saumya', 32, 'Delhi', 30000),
            ('Saumya', 32, 'Mumbai', 20000),
            ('Aaditya', 40, 'Dehradun', 24000),
            ('Seema', 32, 'Delhi', 70000)
            ]

df = pd.DataFrame(employees, columns=['Name', 'Age','City', 'Salary'], )
df = df.set_index('Name')

user_list = ['Age', 'Gender']
data = df[list(set(user_list) & set(df.columns))]
if len(data):
    print(data)
else:
    print('try again')

Answer 2

您需要将user_list 视为一个列表，并检查它是否是df.columns 的子集，在上述实现中，您仅在第一个元素之后决定。

while True:
    user_list = [column for column in input('Please Enter columns you wish to keep: ').split()]
    if set(user_list).issubset(set(df.columns)):
        df = df[user_list]
        break
    else:
        print('try again')
print(df)