0

I am trying to assign all the three unique groups from the group column in df to different variables (see my code) using Python. How do I incorporate this inside a for loop? Obviously var + i does not work.

import pandas as pd

data = {
    'group': ['a', 'a', 'a', 'b', 'b', 'c', 'c'], 
    'num': list(range(7))
    }  
df = pd.DataFrame(data)  

unique_groups = df['group'].unique()

# How do I incorporate this logic inside a for loop??
var1 = df[df['group'] == unique_groups[0]]
var2 = df[df['group'] == unique_groups[1]]
var3 = df[df['group'] == unique_groups[2]]

# My approach:
for i in range(len(unique_groups)):
    var + i = df[df['group'] == unique_groups[i]] # obviously "var + i" does not work
shsh
  • 684
  • 1
  • 7
  • 18
  • 1
    Can your results be an array instead? `var[0], var[1]...`? – Stuart Apr 26 '22 at 16:13
  • You can't easily dynamically assign variables in Python - see https://stackoverflow.com/questions/8028708/dynamically-set-local-variable and https://stackoverflow.com/questions/5036700/how-can-you-dynamically-create-variables – Stuart Apr 26 '22 at 16:14
  • @Stuart Thanks this is the answer I wanted: ```all_vars = []``` ```all_vars.append(df[df['group'] == unique_groups[i]])``` – shsh Apr 26 '22 at 16:24

2 Answers2

1

You can do this using a dictionary, basically:

all_vars ={}
for i in range(len(unique_groups)):

    all_vars[f"var{i}"] = df[df['group'] == unique_groups[i]]
alejandro
  • 11
  • 3
  • Is there any other way without using a dictionary? – shsh Apr 26 '22 at 16:11
  • You could try using globals, check this: https://stackoverflow.com/questions/6181935/how-do-you-create-different-variable-names-while-in-a-loop – alejandro Apr 26 '22 at 16:16
1

From your comment it seems it is okay for all_vars to be a list so that all_vars[0] is the first group, all_vars[1] the second, etc. In that case, consider using groupby instead:

all_vars = [group for name, group in df.groupby("group")]
Stuart
  • 9,597
  • 1
  • 21
  • 30