Condition to groupby or not groupby in Pandas

Question

I'm trying to create a 'conditional groupby' if a certain choice is made from a user input. If the 'ZONE' column exists in the dataframe I want to group by that zone and then iterate over the list of features ['Var1', 'Var2'] .

Except if there is no 'ZONE' column I want to just iterate over the list of features without the groupby.

My pseudo code example is:

import pandas as pd

data = pd.DataFrame({'County' : [1, 2, 2, 2, 3, 3], 'ZONE' : [88, 88, 19, 19, 10, 19], 'Var1' : [78, 90, 97, 100, 12, 140], 'Var2' : [56, 92, 122, 134, 120, 140]})

features = ['Var1', 'Var2']

if 'ZONE' in data.columns:
    data_grouped = data.groupby(['ZONE'])
if 'ZONE' not in data.columns:
    data_grouped = data.copy()

# iterate over grouped zone data
for zone, zone_data in data_grouped:
# iterate over feature columns
      for feature in features:
          data_feature = data_grouped[feature]
          print(data_feature)
          ......make graphs and other things with this grouped data.....

The above code will work for the groupby ZONE case but if there is no ZONE I don't know how to ignore this groupby and iterate only over features in a single for loop - I'd like to have a single for loop for each case instead of breaking both cases and repeating a bunch of graphing code.

Is there any way to do this? Maybe an itertools solution?

Answer 1

I am not sure the what final result do you want. I think there is better way to achieve the final result without loop.

anyway there is simple way to handle what you want:

if 'ZONE' in data.columns:
    data_grouped = data.groupby(['ZONE'])
if 'ZONE' not in data.columns:
    data_grouped = ['NoZone', data]

for zone, zone_data in data_grouped:
# iterate over feature columns
      for feature in features:
          data_feature = zone_data[feature]
          print(data_feature)
          ......make graphs and other things with this grouped data.....