typescript: flatmap implementation

2021-11-18

~1 min read

161 words

Fun little toy problem: implementing a flat map.

The function can take a list of nested arrays of any type and an optional second argument that would apply to each element in the list.

The solution below is O(n) for both space (I built a second array of n length) and time (I iterate over element twice - once to flatten, once to map), where n is the total number of elements.

flatMap.ts

type NestedArray<T> = (NestedArray<T> | T)[]

function flatMap<T>(arr: NestedArray<T>, apply?: (el: T) => T): T[] {
  const flattened: T[] = []

  arr.forEach((el) => {
    if (Array.isArray(el)) {
      flattened.push(...flatMap(el))
    } else {
      flattened.push(el)
    }
  })

  if (apply) {
    return flattened.map(apply)
  }
  return flattened
}

console.log(flatMap([[1, 2, [3]], 4])) // [1, 2, 3, 4]
console.log(flatMap([[1, 2, [3]], 4], (x) => x * 2)) // [2, 4, 6, 8]

The coolest part of this was the opportunity to use a recursive type:

type NestedArray<T> = (NestedArray<T> | T)[]

Hi there and thanks for reading! My name's Stephen. I live in Chicago with my wife, Kate, and dog, Finn. Want more? See about and get in touch!