Heap implementation

Heap is a very useful data structure for solving many programming problems.

One of the problem I found is merge k sorted lists on leetcode.

In this problem we have possible many lists which are sorted and we need to return 1 sorted list, so the most obvious attempt will be keeping always first element of every list in some data structure that allows to very efficient find max or min value (min in this case).

For heap taking min or max is O(1) complexity
And adding new element is O(log n)

This is efficient enough to build O(n log m) algorithm, where n is a collection size of all elements from all lists together and m is number of lists.

Here is my very simple implementation of heap.

const Heap = (valFn = n => n) => {
  const arr = [-1];

  const up = idx => {
    while (idx > 1) {
      const ni = idx / 2 | 0;
      if (valFn(arr[idx]) < valFn(arr[ni])) {
        [arr[idx], arr[ni]] = [arr[ni], arr[idx]];
      }
      idx = ni;
    }
    return idx;
  };

  return {
    add: el => up(arr.push(el) - 1),
    take: () => {
      const len = arr.length;
      if (len <= 1) return [][0];
      let idx = 1;
      const res = arr[idx];
      while (idx < len) {
        const ia = idx * 2;
        const ib = idx * 2 + 1;
        if (ia >= len) break;
        if (ib >= len || valFn(arr[ia]) < valFn(arr[ib])) {
          arr[idx] = arr[ia];
          idx = ia;
        } else {
          arr[idx] = arr[ib];
          idx = ib;
        }
      }
      if (idx === arr.length - 1) {
        arr.pop();
      } else {
        arr[idx] = arr.pop();
        up(idx);
      }
      return res;
    },
    size: () => arr.length - 1,
    data: () => arr.slice(1),
  };
};

Be noticed that constructor of this heap is expecting getter function for element value.
Thanks to that we can keep in our data structure more information.
For this problem we will need to keep index number of list so we can take next element from same list if we removing it from heap.

Here is my very simple implementation of this problem using arrays instead lists and using my heap implementation.

const mergeKLists = lists => {
  const heap = Heap(a => a[0]);
  lists.map((list, i) => {
    const el = list.shift();
    heap.add([el, i]);
  });
  const res = [];
  while (1) {
    const next = heap.take();
    if (next[0] === [][0]) break;
    const [val, i] = next;
    heap.add([lists[i].shift(), i]);
    res.push(val);
  }
  return res;
};

const lists = [[1,4,5],[1,3,4],[2,6]];

mergeKLists(lists);

You can play with this code on Instacode.app

Changing array to list is very trivial thing.

I also added this heap to my utils collection for future tasks.

Be noticed that I implemented min heap here actually.
In most cases it is good enough because even if you need max heap instead and you don't care about handling same elements like 1 === 1 you can just pass invert value getter function like this below to accomplish that.

const heap = Heap(n => Number.MAX_SAFE_INTEGER - n);