Search code examples
c#linqbinning

Split a list into n equal parts


Given a sorted list, and a variable n, I want to break up the list into n parts. With n = 3, I expect three lists, with the last one taking on the overflow.

I expect: 0,1,2,3,4,5, 6,7,8,9,10,11, 12,13,14,15,16,17

If the number of items in the list is not divisible by n, then just put the overflow (mod n) in the last list.

This doesn't work:

static class Program
{
    static void Main(string[] args)
    {
        var input = new List<double>();
        for (int k = 0; k < 18; ++k)
        {
            input.Add(k);
        }
        var result = input.Split(3);
        
        foreach (var resul in result)
        {
            foreach (var res in resul)
            {
                Console.WriteLine(res);
            }
        }
    }
}

static class LinqExtensions
{
    public static IEnumerable<IEnumerable<T>> Split<T>(this IEnumerable<T> list, int parts)
    {
        int i = 0;
        var splits = from item in list
                     group item by i++ % parts into part
                     select part.AsEnumerable();
        return splits;
    }
}

Solution

  • I think you would benefit from Linq's .Chunk() method.

    If you first calculate how many parts will contain the equal item count, you can chunk list and yield return each chunk, before yield returning the remaining part of list (if list is not divisible by n).

    As pointed out by Enigmativity, list should be materialized as an ICollection<T> to avoid possible multiple enumeration. The materialization can be obtained by trying to cast list to an ICollection<T>, and falling back to calling list.ToList() if that is unsuccessful.

    A possible implementation of your extension method is hence:

    public static IEnumerable<IEnumerable<T>> Split<T>(this IEnumerable<T> list, int parts)
    {
        var collection = list is ICollection<T> c
            ? c
            : list.ToList();
        
        var itemCount = collection.Count;
        
        // return all items if source list is too short to split up
        if (itemCount < parts)
        {
            yield return collection;
            yield break;
        }
        
        var itemsInEachChunk = itemCount / parts;
        
        var chunks = itemCount % parts == 0
            ? parts
            : parts - 1;
        
        var itemsToChunk = chunks * itemsInEachChunk;
        
        foreach (var chunk in collection.Take(itemsToChunk).Chunk(itemsInEachChunk))
        {
            yield return chunk;
        }
        
        if (itemsToChunk < itemCount)
        {
            yield return collection.Skip(itemsToChunk);
        }
    }
    

    Example fiddle here.