LeetCode 2631 - Group By

Difficulty: 🟡 Medium
Topics: —

Solution

Problem Understanding

The problem asks us to enhance arrays such that any array can call a groupBy method with a callback function fn. This function will determine the key for grouping each element. The output should be an object (dictionary in Python, map in Go) where each key corresponds to the result of fn applied to an array element, and the value is a list of all elements that map to that key. The ordering of elements within each group must preserve their original order in the array, although the order of keys in the result does not matter.

The input can be heterogeneous: arrays of objects, arrays of arrays, or arrays of primitive values. The callback always returns a string key. The array size can range up to 10^5, so we need a solution that scales linearly in time and does not use excessive additional memory.

Important edge cases include an empty array, arrays with all elements mapping to the same key, arrays where each element maps to a unique key, and arrays with nested structures or non-hashable types in other languages.

Approaches

The brute-force approach would involve iterating over the array and, for each element, checking whether its group key already exists in the result object. If it does, append the element to that group; otherwise, create a new group. This approach is simple and works for small inputs, but it still scales linearly in time, which is acceptable here. The key insight is that using a hash map allows constant-time access to existing groups, so the overall complexity remains O(n) without nested loops.

The optimal solution leverages the hash map (or dictionary) to accumulate elements per group key as we traverse the array once. This guarantees both efficiency and correctness while maintaining element order within groups.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n)	O(n)	Iterate array, check and append to object groups; works linearly with hash map
Optimal	O(n)	O(n)	Single-pass accumulation into a dictionary or map using key from `fn`

Algorithm Walkthrough

Initialize an empty dictionary (or map) result to hold grouped elements.
Iterate over each element item in the input array.
Apply the callback function fn to item to get the group key.
If the key does not exist in result, create a new empty list for that key.
Append item to the list corresponding to its key in result.
After processing all elements, return the result object.

Why it works: The dictionary ensures each key accumulates all elements that match it. Iterating once guarantees O(n) time. Preserving insertion order within lists satisfies the problem requirement.

Python Solution

from typing import List, Callable, Any, Dict

class ArrayWithGroupBy(list):
    def groupBy(self, fn: Callable[[Any], str]) -> Dict[str, List[Any]]:
        result: Dict[str, List[Any]] = {}
        for item in self:
            key = fn(item)
            if key not in result:
                result[key] = []
            result[key].append(item)
        return result

The implementation extends Python's list class. The groupBy method creates a dictionary to collect groups. Each element is processed in order, and its key is computed using fn. If the key does not exist in the dictionary, it initializes a new list; otherwise, it appends to the existing list. Finally, the dictionary is returned.

Go Solution

package main

type Any interface{}

func GroupBy(array []Any, fn func(Any) string) map[string][]Any {
    result := make(map[string][]Any)
    for _, item := range array {
        key := fn(item)
        result[key] = append(result[key], item)
    }
    return result
}

In Go, we define a GroupBy function that accepts a slice of empty interface Any and a callback function. A map is used to collect the groups. The append function naturally handles initialization of slices, so we do not need to explicitly check for key existence. This preserves element order within each group.

Worked Examples

Example 1:

Input: [{"id":"1"},{"id":"1"},{"id":"2"}], fn = item -> item.id

Iteration	Item	Key	Result
1	{"id":"1"}	"1"	{"1": [{"id":"1"}]}
2	{"id":"1"}	"1"	{"1": [{"id":"1"},{"id":"1"}]}
3	{"id":"2"}	"2"	{"1": [...], "2": [{"id":"2"}]}

Example 2:

Input: [[1,2,3],[1,3,5],[1,5,9]], fn = list -> String(list[0])

Iteration	Item	Key	Result
1	[1,2,3]	"1"	{"1": [[1,2,3]]}
2	[1,3,5]	"1"	{"1": [[1,2,3],[1,3,5]]}
3	[1,5,9]	"1"	{"1": [[1,2,3],[1,3,5],[1,5,9]]}

Example 3:

Input: [1,2,3,4,5,6,7,8,9,10], fn = n -> String(n>5)

Iteration	Item	Key	Result
1-5	1-5	"false"	"false": [1,2,3,4,5]
6-10	6-10	"true"	"true": [6,7,8,9,10]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each element is visited exactly once and key lookup in dictionary/map is O(1) on average.
Space	O(n)	Additional dictionary/map stores all elements grouped by key, same size as input in worst case.

This ensures efficiency even for arrays with up to 10^5 elements.

Test Cases

# Empty array
assert ArrayWithGroupBy([]).groupBy(lambda x: str(x)) == {}

# All elements same key
arr = ArrayWithGroupBy([1,1,1])
assert arr.groupBy(lambda x: "same") == {"same": [1,1,1]}

# Unique keys
arr = ArrayWithGroupBy([1,2,3])
assert arr.groupBy(lambda x: str(x)) == {"1":[1],"2":[2],"3":[3]}

# Example 1
arr = ArrayWithGroupBy([{"id":"1"},{"id":"1"},{"id":"2"}])
assert arr.groupBy(lambda x: x["id"]) == {"1":[{"id":"1"},{"id":"1"}],"2":[{"id":"2"}]}

# Example 2
arr = ArrayWithGroupBy([[1,2,3],[1,3,5],[1,5,9]])
assert arr.groupBy(lambda x: str(x[0])) == {"1":[[1,2,3],[1,3,5],[1,5,9]]}

# Example 3
arr = ArrayWithGroupBy([1,2,3,4,5,6,7,8,9,10])
assert arr.groupBy(lambda x: str(x>5)) == {"true":[6,7,8,9,10],"false":[1,2,3,4,5]}

Test	Why
Empty array	Tests handling zero elements
All elements same key	Ensures grouping works when everything maps to one key
Unique keys	Ensures algorithm handles maximum number of groups
Example 1	Verifies object grouping by property
Example 2	Verifies array-of-array grouping
Example 3	Verifies numeric condition-based grouping

Edge Cases

Empty array: An array with no elements should return an empty dictionary. This is handled naturally because no iteration occurs, and the initial empty dictionary is returned.

All elements map to same key: The algorithm correctly appends each element to the same list in the dictionary, preserving order.

Heterogeneous types: Arrays may contain objects, arrays, or primitives. The algorithm relies on the callback fn to compute a string key, which guarantees compatibility with diverse element types. As long as fn returns a string, grouping works regardless of element type.