Problem

Given a function fn, return a memoized version of that function.

A **memoized **function is a function that will never be called twice with the same inputs. Instead it will return a cached value.

fn can be any function and there are no constraints on what type of values it accepts. Inputs are considered identical if they are === to each other.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

Input: 
getInputs = () => [[2,2],[2,2],[1,2]]
fn = function (a, b) { return a + b; }
Output: [{"val":4,"calls":1},{"val":4,"calls":1},{"val":3,"calls":2}]
Explanation:
const inputs = getInputs();
const memoized = memoize(fn);
for (const arr of inputs) {
  memoized(...arr);
}

For the inputs of (2, 2): 2 + 2 = 4, and it required a call to fn().
For the inputs of (2, 2): 2 + 2 = 4, but those inputs were seen before so no call to fn() was required.
For the inputs of (1, 2): 1 + 2 = 3, and it required another call to fn() for a total of 2.

Example 2

1
2
3
4
5
6

Input: 
getInputs = () => [[{},{}],[{},{}],[{},{}]] 
fn = function (a, b) { return ({...a, ...b}); }
Output: [{"val":{},"calls":1},{"val":{},"calls":2},{"val":{},"calls":3}]
Explanation:
Merging two empty objects will always result in an empty object. It may seem like there should only be 1 call to fn() because of cache-hits, however none of those objects are === to each other.

Example 3

1
2
3
4
5
6

Input: 
getInputs = () => { const o = {}; return [[o,o],[o,o],[o,o]]; }
fn = function (a, b) { return ({...a, ...b}); }
Output: [{"val":{},"calls":1},{"val":{},"calls":1},{"val":{},"calls":1}]
Explanation:
Merging two empty objects will always result in an empty object. The 2nd and 3rd third function calls result in a cache-hit. This is because every object passed in is identical.

Constraints

  • 1 <= inputs.length <= 10^5
  • 0 <= inputs.flat().length <= 10^5
  • inputs[i][j] != NaN

Solution

Method 1 – Closure with WeakMap for Argument Identity

Intuition

To memoize a function for any type of arguments, we need to cache results based on argument identity (===). For primitive values, we can use a Map; for objects, we use a WeakMap to avoid memory leaks. We build a nested cache structure, where each argument leads to a new Map/WeakMap, and the result is stored at the final node.

Approach

  1. Create a root cache (Map).
  2. For each call, traverse the cache using each argument:
    • Use WeakMap for objects/functions, Map for primitives.
    • At each level, create a new Map/WeakMap if not present.
  3. At the final level, store or retrieve the result.
  4. Track the number of times the original function is called (for getCallCount).

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

function memoize(fn) {
    let callCount = 0;
    const cache = new Map();
    function getCache(arg, map) {
        if (typeof arg === 'object' && arg !== null || typeof arg === 'function') {
            if (!map.has(arg)) map.set(arg, new WeakMap());
            return map.get(arg);
        } else {
            if (!map.has(arg)) map.set(arg, new Map());
            return map.get(arg);
        }
    }
    const memoized = (...args) => {
        let node = cache;
        for (let i = 0; i < args.length; ++i) {
            node = getCache(args[i], node);
        }
        if (node.has('result')) return node.get('result');
        callCount++;
        const res = fn(...args);
        node.set('result', res);
        return res;
    };
    memoized.getCallCount = () => callCount;
    return memoized;
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

function memoize<T extends (...args: any[]) => any>(fn: T): T & { getCallCount: () => number } {
    let callCount = 0;
    const cache = new Map<any, any>();
    function getCache(arg: any, map: Map<any, any>): any {
        if ((typeof arg === 'object' && arg !== null) || typeof arg === 'function') {
            if (!map.has(arg)) map.set(arg, new WeakMap());
            return map.get(arg);
        } else {
            if (!map.has(arg)) map.set(arg, new Map());
            return map.get(arg);
        }
    }
    const memoized = (...args: Parameters<T>): ReturnType<T> => {
        let node = cache;
        for (let i = 0; i < args.length; ++i) {
            node = getCache(args[i], node);
        }
        if (node.has('result')) return node.get('result');
        callCount++;
        const res = fn(...args);
        node.set('result', res);
        return res;
    };
    (memoized as any).getCallCount = () => callCount;
    return memoized as T & { getCallCount: () => number };
}

Complexity

  • ⏰ Time complexity: O(k) for each call, where k is the number of arguments.
  • 🧺 Space complexity: O(n), where n is the number of unique argument sets cached.