how to split dataframe into equal number of subset in python - python-3.x

Question

I have a dataframe
import pandas as pd
d = {'user': [1, 1, 2,2,2,2 ,2,2,2,2], 'friends':
[1,2,1,5,4,6,7,20,9,7]}
df = pd.DataFrame(data=d)
I try to split the df into several n pieces in a loop. For example, for n=3
n=3
for i in range(3):
subdata = dosomething(df)
print(subdata)
the output will be someting like
# first loop
user friends
0 1 1
1 1 2
2 2 1
3 2 5
# second loop
user friends
0 2 4
1 2 6
2 2 7
3 2 20
#third loop
user friends
0 2 9
1 2 7

Answer 1

You can use iloc and loop through the dataframe, put each new dataframe in a dictionary for recall later.
dfs = {}
chunk = 4
Loop through the dataframe by chunk sizes. Create df and add to dict.
for n in range((df.shape[0] // chunk + 1)):
df_temp = df.iloc[n*chunk:(n+1)*chunk]
df_temp = df_temp.reset_index(drop=True)
dfs[n] = df_temp
Use this if statement for any left over rows at the end.
if df.shape[0] % chunk != 0:
df_temp = df.iloc[-int(df.shape[0] % chunk):]
df_temp = df_temp.reset_index(drop=True)
dfs[n] = df_temp
else:
pass
Access the dataframes in the dictionary.
print(dfs[0])
user friends
0 1 1
1 1 2
2 2 1
3 2 5
print(dfs[1])
user friends
0 2 4
1 2 6
2 2 7
3 2 20
print(dfs[2])
user friends
0 2 9
1 2 7