def weight_imbalance(sizes, i):
    """Calculate the weight imbalance between two subgroups separated at index i"""
    sizes = list(sizes)
    if i == 0 or i == len(sizes) - 1 or i == len(sizes):
        return sum(sizes)
    elif 0 < i < len(sizes) - 1:
        return abs(sum(sizes[:i]) - sum(sizes[i:]))
    else:
        raise IndexError(
            f'Split index {i} out of range for sizes with len of {len(sizes)}')
