CODEWITHSUNDEEP
pandas
fred.schwartz
fred.schwartz

Reputation: 2155

Pandas get list of columns if columns name contains

I have written this code to show a list of column names in a dataframe if they contains 'a','b' ,'c' or 'd'.

I then want to say trim the first 3 character of the column name for these columns.

However its showing an error. Is there something wrong with the code?

ind_cols= [x for x in df if df.columns[df.columns.str.contains('|'.join(['a','b','c','d']))]]

df[ind_cols].columns=df[ind_cols].columns.str[3:]

Upvotes: 2

Views: 161

Answers (1)

jezrael
jezrael

Reputation: 863246

Use list comprehension with if-else:

L = df.columns[df.columns.str.contains('|'.join(['a','b','c','d']))]
df.columns = [x[3:] if x in L else x for x in df.columns]

Another solution with numpy.where by boolean mask:

m = df.columns.str.contains('|'.join(['a','b','c','d']))
df.columns = np.where(m, df.columns.str[3:], df.columns)

Upvotes: 2

Related Questions