Split Pandas dataset column based on values (suffixes: string operation)

Question

In Python using Pandas, I am splitting a dataset column into 4 lists based on the suffix of the values. For the 3 suffixes I am using a list comprehension then for the 4th one, a set operation that substracts the 3 lists from the original list with all values:

import pandas as pd df = pd.DataFrame({ "alcohol_by_volume": [], "barcode": [], "calcium_per_hundred": [], "calcium_unit": [], "carbohydrates_per_hundred": [], "carbohydrates_per_portion": [], "carbohydrates_unit": [], "cholesterol_per_hundred": [], "cholesterol_unit": [], "copper_cu_per_hundred": [], "copper_cu_unit": [], "country": [], "created_at": [], "energy_kcal_per_hundred": [], "energy_kcal_per_portion": [], "energy_kcal_unit": [], "energy_per_hundred": [], "energy_per_portion": [], "energy_unit": [], "fat_per_hundred": [], "fat_per_portion": [], "fat_unit": [], "fatty_acids_total_saturated_per_hundred": [], "fatty_acids_total_saturated_unit": [], "fatty_acids_total_trans_per_hundred": [], "fatty_acids_total_trans_unit": [], "fiber_insoluble_per_hundred": [], "fiber_insoluble_unit": [], "fiber_per_hundred": [], "fiber_per_portion": [], "fiber_soluble_per_hundred": [], "fiber_soluble_unit": [], "fiber_unit": [], "folate_total_per_hundred": [], "folate_total_unit": [], "folic_acid_per_hundred": [], "folic_acid_unit": [], "hundred_unit": [], "id": [], "ingredients_en": [], "iron_per_hundred": [], "iron_unit": [], "magnesium_per_hundred": [], "magnesium_unit": [], "manganese_mn_per_hundred": [] }) colnames_all = df.columns.to_list() colnames_unit = [n for n in colnames_all if n.endswith("_unit")] colnames_per_hundred = [n for n in colnames_all if n.endswith("_per_hundred")] colnames_per_portion = [n for n in colnames_all if n.endswith("_per_portion")] colnames_other = list( set(colnames_all) - set(colnames_unit + colnames_per_hundred + colnames_per_portion) ) 

Expected result (2 examples, other 2 lists are similar to 1st one):

colnames_unit: ['calcium_unit', 'carbohydrates_unit', 'cholesterol_unit', 'copper_cu_unit', 'energy_kcal_unit', 'energy_unit', 'fat_unit', 'fatty_acids_total_saturated_unit', 'fatty_acids_total_trans_unit', 'fiber_insoluble_unit', 'fiber_soluble_unit', 'fiber_unit', 'folate_total_unit', 'folic_acid_unit', 'hundred_unit', 'iron_unit', 'magnesium_unit'] colnames_other: ['ingredients_en', 'country', 'id', 'created_at', 'barcode', 'alcohol_by_volume'] 

However this does not look like the best way to do this. Is there a "better" way, i.e. shorter and/or more elegant/idiomatic?

Answer 1

colnames_all = df.columns.to_list() 

I don't see a clear need for this. We could simply refer to df.columns instead.

---

list( set(colnames_all) - set(colnames_unit + colnames_per_hundred + colnames_per_portion) ) 

That doesn't seem so bad, to me. Certainly the intent is clear.

---

colnames_unit = [n for n in colnames_all if n.endswith("_unit")] 

Consider rephrasing this as

colnames_unit = [n for n in colnames_all if re.search(r'_unit$', n)] 

That lets us generalize in this way:

colnames_measured = [n for n in df.columns if re.search(r'_(unit|per_hundred|per_portion)$', n)] 

To find the inverse:

colnames_other = [n for n in df.columns if not re.search(r'_(unit|per_hundred|per_portion)$', n)] 

Answer 2

Don't use comprehensions. Don't use lists. Don't use sets. Use Pandas string vectorisation:

colnames_all = df.columns is_unit = colnames_all.str.endswith("_unit") is_hundred = colnames_all.str.endswith("_per_hundred") is_portion = colnames_all.str.endswith("_per_portion") colnames_unit = colnames_all[is_unit] colnames_per_hundred = colnames_all[is_hundred] colnames_per_portion = colnames_all[is_portion] colnames_other = colnames_all[~(is_unit | is_hundred | is_portion)] print(colnames_other) 

Index(['alcohol_by_volume', 'barcode', 'country', 'created_at', 'id', 'ingredients_en'], dtype='object')