How to remove columns based on user input in python (pandas)?

Question

How to write code so that the dataframe remove column based on user input of the column name?

amend_col= input(" would u like to 'remove' any column?:")
if amend_col == "yes":
   s_col= input("Select the column you want to add or remove:")
   if s_col in df.columns:
      print("Column is found and removed")
      df.drop(df[s_col])
else:
    print("No columns removed")
            

Answer 1

You can amend your codes as follows:

amend_col= input(" would u like to 'remove' any column?:")
if amend_col == "yes":
   s_col= input("Select the column you want to add or remove:")
   if s_col in df.columns and amend_col =="yes":
      print("Column is found and removed")
      df = df.drop(columns=s_col)
else:
    print("No columns removed")

You code is close, we just need to replace the line df.drop(df[s_col]) with

df = df.drop(columns=s_col)