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")

Hello, I think your code should work, except that if you want to modify df you have to add `inplace=True` to the drop method, and specify `s_col` directly instead of `df[s_col]`. — totooooo, May 17 '21 at 17:00
@totooooo Using `inplace=True` is not recommended and is seriously considered to be deprecated in future versions. See [this post](https://stackoverflow.com/a/60020384/15070697) and also [this article](https://www.dataschool.io/future-of-pandas/#inplace) — SeaBean, May 17 '21 at 19:00

score 2 · Accepted Answer · answered May 17 '21 at 16:34

2

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)

answered May 17 '21 at 16:34

SeaBean

22,547
3
13
25

I was lazy to write so I copy your solution(95%) lol....btw sorry for that mate.....I want to delete that as It a copy of your solution but my deletion was blocked yesterday...again sorry for that and also please don't embrassed me by showing your screen dump that you didn't downvote **:)** – Anurag Dabas May 18 '21 at 02:43

score 0 · Answer 2 · answered May 17 '21 at 16:41

amend_col= input("Would u like to 'remove' any column?:")
# Maybe check case or input type
if amend_col == "yes":
   s_col= input("Select the column you want to add or remove:")
   # Again maybe check case, put all the col ids to uppercase
   try:
      df.drop([s_col],axis=1)
      print("Column {} has been found and removed.".format(s_col))
   except:
      print("Column {} does not exist.\nNo columns have been removed.".format(s_col))
else:
    print("No columns removed.")

You can move your code in to a try, except replacing if s_col in df.columns and amend_col =="yes": .... Firstly, you dont need to check and amend_col =="yes" as it was checked in the previous if statement. Secondly, the try, except performs the same function and provides feedback.

You could also do it with an if:

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([s_col],axis=1)
   else:
      print("Column was not found")
else:
    print("No columns removed")

score 0 · Answer 3 · answered May 17 '21 at 16:44

You don't need to write again amend_col == 'yes' in the second statement since you have already pass that first one.

Also you need to specify on what axis you want to do the drop, and in your case is axis = 1, since you want to drop a column.

So your code will be something like this...

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 = df.drop(s_col, axis=1)
    else:
        print("Column is not in the dataframe."
else:
    print("No columns removed")

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

3 Answers3