How to apply panda group by with minimum and maximum size condition (Pythonic way)

Question

I have a dataframe in pandas which I need to group and store in a new array where I need the size of every group with a specific size and if one exceeds the minimum size, it should be added to one of the previous groups that have the smallest size. For example, after I grouped the data, I will have groups G that are len(G)<=b, len(G)>=a, or a <= len(G) <= b. So, I need to make the groups with len(G)>=a to meet the condition a <= len(G) <= b.

The code is working now. So, I would like to know if there is a more convenient way to do that.

import numpy as np
import pandas as pd

rng = np.random.default_rng()  # Just for testing
df = pd.DataFrame(rng.integers(0, 10, size=(1000, 4)), columns=list('ABCD'))
# The dataframe is grouped depend on specific column.
ans = [pd.DataFrame(y) for x, y in df.groupby(df.columns[3], as_index=False)] 

n = 20 # The maximum size of the group is 25

new_arrayi_index = 0
new_array = []
for count_index in range(len(ans)):
    l = ans[count_index]
   
    if len(l) > n:

        df_shuffled = pd.DataFrame(l).sample(frac=1)
        final = [df_shuffled[i:i+n] for i in range(0,df_shuffled.shape[0],n)]

        for inde in range(len(final)):
            if len(final[inde]) <= 5 and new_arrayi_index != 0: #The minimum size of the group is 5

                new_array[new_arrayi_index - 1]=new_array[new_arrayi_index - 1]+final[inde]

            else:
                new_array.append(final[inde])
                new_arrayi_index += 1

    else: 

        new_array.append(l)
        new_arrayi_index += 1

count_index_ = 0
for count_index in range(len(new_array)):
    print("count", count_index, "Size", len(new_array[count_index]))
    print(new_array[count_index])
    count_index_ += count_index

print(count_index_)

Answer 1

I was following this post since the beginning curious about how the discussion would g, because the OP's problem is not always possible to solve.

Existence of a solution

Take the following example: A group has 19 elements and you want to split it in sections of size between 10 and 15.

The solution exists if and only if exists an integer g, such that n/b <= g <= n/a. In this case you can see that g sections of length a will use g*a <= n elements, and sections of length b will use g*b >= n.

In this situation it is also possible to have a balanced partition, in the sense that the largest section will be at most one record larger than the smallest section (the smallest will have n//g records).

Restating the problem

We could do a slight modification to the problem as split in the minimum possible number of sections containing at most b records each. Such that the length of each section satisfy a <= len(s) <= a+1.

Notice that in this case we are adjusting a to be the closest possible from b so that the problem will have a solution. For solvable problems the solution will be a solution to the original problem, for problems that can't be solved it will modify the original requirement by reducing a so that the problem can be solved.

The example above would become: Split 19 elements in the minimum possible number of balanced groups with no more than 15 elements. Then the solution is having one section of 10 and one section of 9 elements.

A pythonic solution

def group_and_split(df, b, column):
    '''
    - df    : a datafame
    - b     : the largest allowed section
    - column: the column by which the data must be grouped
    '''
    
    # doing it in a pythonic way
    return [[y] if len(y) <= n else np.array_split(y, (len(y)+b-1)//b)
             for x, y in df.groupby(column, as_index=False)]

You can check that it gives a solution to the restated problem

pd.DataFrame([{
    'num-sections': len(g), 
    'largest-section': max(len(gi) for gi in g), 
    'smallest-sections':min(len(gi) for gi in g)
} for g in group_and_split(df, 25, 'D')])

A complete running code

import pandas as pd
import numpy as np

rng = np.random.default_rng(1)  # Just for testing
df = pd.DataFrame(rng.integers(0, 10, size=(1000, 4)), columns=list('ABCD'))

def group_and_split(df, b, column, n=1):
    '''
    - df    : a datafame
    - b     : the largest allowed section
    - column: the column by which the data must be grouped
    - n     : don't split groups that are smaller than this
    '''

    # doing it in a pythonic way
    return [[y] if len(y) <= n else np.array_split(y, (len(y)+b-1)//b)
             for x, y in df.groupby(column, as_index=False)]



pd.DataFrame([{
    'num-sections': len(g), 
    'largest-section': max(len(gi) for gi in g), 
    'smallest-sections':min(len(gi) for gi in g)
} for g in group_and_split(df, 25, 'D', 93)])

Printing all the groups

for group in group_and_split(df, 25, 'D'):
    for section in group:
        print(section)

Answer 2

I wrote a function that splits the dataframe into chunks that are equal to the max size. It checks the size of the remainder for the last chunk, and if the remainder is smaller than the minimum size, it splits the last two chunks into two chunks of approximately equal size.

Building off answer at Split a large pandas dataframe

import numpy as np
import pandas as pd


rng = np.random.default_rng(seed=1)  # Just for testing
df = pd.DataFrame(rng.integers(0, 10, size=(1000, 4)), columns=list('ABCD'))
# The dataframe is grouped depend on specific column.

n = 20  # The maximum size of the group is 25


# https://stackoverflow.com/questions/17315737/split-a-large-pandas-dataframe

def split_dataframe(df, chunk_size=20, min_size=10):

    chunks = list()
    remainder = len(df) % chunk_size

    if 0 < remainder < min_size:
        num_chunks = len(df) // chunk_size - 1
        for i in range(num_chunks):
            chunks.append(df[i * chunk_size:(i + 1) * chunk_size])
        df_ = df[(num_chunks) * chunk_size:]
        last_break = int(len(df_) / 2)
        chunks.append(df_[:last_break])
        chunks.append(df_[last_break:])
        return chunks
    else:
        num_chunks = len(df) // chunk_size + 1
        if remainder == 0:
            num_chunks += -1
        for i in range(num_chunks):
            chunks.append(df[i*chunk_size:(i+1)*chunk_size])
        return chunks


new_array = []
for group, df_ in df.groupby(df.columns[3], as_index=False):
    if len(df_) > n:
        new_array.extend(split_dataframe(df_))
    else:
        new_array.extend(df_)

count_index_ = 0
for count_index in range(len(new_array)):
    print("count", count_index, "Size", len(new_array[count_index]))
    print(new_array[count_index])
    count_index_ += count_index

print(count_index_)

Answer 3

change this line -> ans = [pd.DataFrame(y) for x, y in df.groupby(df.columns[3], as_index=False)] to ans = [pd.DataFrame(y) for x, y in df.groupby(df.columns[3].min(), as_index=False)] for min

and ans = [pd.DataFrame(y) for x, y in df.groupby(df.columns[3].max(), as_index=False)] for max