Recreating lodash partition method

Lodash _.partition function splits an array into two groups one filled with the items fullfiling the provided condition, and the other group is filled with the items that don’t.

The goal of this blog article is to duplicate the partition method, but with some modifications and add some extra features. the concept will remain the same, however instead of taking one predicate, our function will be able to take an array of predicates (partition functions) and partition a giver array based on them.

Our function signature in Typescript will look like this

type PartitionWith = <T>(items: T[], predicates: ((item: T) => boolean)[]): T[][]

An example of usage would be partitioning an array into two arrays, one containing number greater than 5, and another with items less or equal than five.

const array = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
const isLessOrEqualThanFive = (number: number) => number <= 5; 
const isGreaterThanFive = (number) => number > 5;

const results = partitionWith(array, [isLessOrEqualThanFive, isGreaterThanFive ]); 

console.log(results); // [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 10]]

We can see that the partition counts are equal to the length of the predicates array, let's write our first assertion and then implement the code to make it pass.

it('creates an array of partitions with a length that is equal to the predicates array length', () => {
        const predicateOne = (n: number) => n < 5; 
        const predicateTwo = (n: number) => n >= 5; 

        const array = [1, 2, 4, 5, 6]; 

        const results = partitionWith(array, [predicateOne, predicateTwo]); 

        expect(results.length).toBe(2); 
})

Passing two predicates, means the resulting array should also contains two partitions.

const partitionWith: PartitionWith = <T>(items: T[], predicates: ((item: T) => boolean)[]) => {
    const results: T[][] = [...Array(predicates.length)].map(x => []); 

    return results; 
}

Our function now creates an array of arrays with a length equal to the predicates array length.

The next step would be implementing the predicates applying logic, the idea is whenever a predicate returns true for an item, the latter will be added to the partitions array at that predicate index.

To find the predicate index, we will use .findIndex function that returns either the index of the first item that satisify the provided condition, or -1 when none is found.

const predicateIndex = predicates.findIndex(predicate => predicate(item));

Let’s write a test before implementing the functionality.

it('create partitions based on the provided predicates', () => {
        const arrayToPartition = [0, 1, '1', 2, 3, 4, '12', 5, 6, 7, 8, 9, , '11', 10]; 

        const isLessThanFive = (maybeNumber: number | string) => typeof maybeNumber === 'number' && maybeNumber < 5; 
        const isGreaterOrEqualThanFive = (maybeNumber: number | string) => typeof maybeNumber === 'number' &&  maybeNumber >= 5;
        const isString = (maybeString: number | string) => typeof maybeString === 'string';  

        const results = partitionWith(arrayToPartition, [isLessThanFive, isGreaterOrEqualThanFive, isString]); 

        expect(results).toEqual([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 10], ['1', '12', '11']]); 
    });

type PartitionWith =  <T>(items: T[], predicates: ((item: T) => boolean)[]) => T[][];
export const partitionWith: PartitionWith = <T>(items: T[], predicates: ((item: T) => boolean)[]) => {
    const results: T[][] = [...Array(predicates.length)].map(x => []); 

    items.forEach((item) => {
        const predicateIndex = predicates.findIndex(predicate => predicate(item)); 

        if(predicateIndex !== -1) {
            results[predicateIndex].push(item); 
        }
    })

    return results; 
}

As stated before, for each element, we try to find which predicate it fullfils, if found we add it to the corresponding predicate index using results[predicateIndex].push(item);.

The solution now ignores all the items that do not satisify any condition. However, the original lodash _.partition function splits the array into two groups, one contains elements that satisify the condition and the second contains elements that do not.

So let's implement that, but first as usual we will write the test before implementing the logic.

it('returns an extra array of items that did not satisfy any condition', () => {
        const items = [0, 1, '1', 2, 3, 4, '12', 5, 6, 7, 8, 9, , '11', 10]; 

        const isLessThanFive = (maybeNumber: number | string) => typeof maybeNumber === 'number' && maybeNumber < 5; 
        const isGreaterOrEqualThanFive = (maybeNumber: number | string) => typeof maybeNumber === 'number' &&  maybeNumber >= 5;

        const results = partitionWith(items, [isLessThanFive, isGreaterOrEqualThanFive]); 
        expect(results).toEqual([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 10], ['1', '12', '11']])

    })

Here we have 2 conditions, we will get only the items that are numbers and less than five, or greater than five, the items left should be added to an array at the end of the partitions array. So let’s first add this array that will contains the falsy items to the results array.

results.push([])

Whenever an item fails to satisfy at least one of the given predicates, it will be added the array which index is at end of the results array. The algorithm will look like this: We should also refactor our previous tests accordignly.

export const partitionWith: PartitionWith = <T>(items: T[], predicates: ((item: T) => boolean)[]) => {
    const results: T[][] = [...Array(predicates.length)].map(x => []); 

    results.push([])

    items.forEach((item) => {
        const predicateIndex = predicates.findIndex(predicate => predicate(item)); 

        if(predicateIndex !== -1) {
            results[predicateIndex].push(item); 
        } else {
            const falsyResultsArrayIndex = predicates.length; 

            results[falsyResultsArrayIndex].push(item);
        }
    })

    return results; 
}

Since the falsy items array is added at the end of the results array, it’s index will be the predicates.length.

Now our partitionWith function will behave exactly the same as partition from lodash when provided only one predicate, keeping the falsy elements into a seperate array at the tail of the results array.

The full version with the tests suite can be found here

Follow me on twitter for more
Ab_del