Top K Frequent Elements

Description

Given a non-empty array of integers, return the k most frequent elements.

Example 1:

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

Example 2:

Input: nums = [1], k = 1
Output: [1]

Note:

• You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
• Your algorithm's time complexity must be better than O(n log n), where n is the array's size.
• It's guaranteed that the answer is unique, in other words the set of the top k frequent elements is unique.
• You can return the answer in any order.

Solution(javascript)

const swap = (a, b, arr) => { // eslint-disable-line
  if (a !== b) {
    const temp = arr[a]
    arr[a] = arr[b] // eslint-disable-line
    arr[b] = temp // eslint-disable-line
  }
}

const Heap = compareFn => (arr = []) => {
  const left = index => 2 * index + 1
  const right = index => 2 * index + 2
  const parent = index => Math.floor((index - 1) / 2)
  const size = () => arr.length

  // log(n)
  const heapify = (index) => {
    const l = left(index)
    const r = right(index)
    let current = index
    if (compareFn(arr[current], arr[l]) > 0 && (l < size())) {
      current = l
    }
    if (compareFn(arr[current], arr[r]) > 0 && (r < size())) {
      current = r
    }
    if (current !== index) {
      swap(current, index, arr)
      heapify(current)
    }
  }
  // log(n)
  const heapifyUp = (index) => {
    const p = parent(index)
    if (p >= 0 && compareFn(arr[p], arr[index]) > 0) {
      swap(p, index, arr)
      heapifyUp(p)
    }
  }
  // O(n)
  const buildHeap = () => {
    for (let i = Math.floor(arr.length / 2); i >= 0; i--) {
      heapify(i)
    }
  }
  const extract = () => {
    swap(0, arr.length - 1, arr)
    const top = arr.pop()
    heapify(0)
    return top
  }
  const remove = (item) => {
    const index = arr.findIndex(x => compareFn(x, item) === 0)
    if (index === -1) {
      return
    }
    arr[index] = arr.pop() // eslint-disable-line
    const p = parent(index)
    if (p < 0 || compareFn(p, arr[index]) < 0) {
      heapify(index)
    } else {
      heapifyUp(index)
    }
  }
  buildHeap()
  return {
    getHeap: () => arr,
    peek: () => {
      if (arr.length === 0) {
        return null
      }
      return arr[0]
    },
    add: (item) => {
      arr.push(item)
      heapifyUp(arr.length - 1)
    },
    extract,
    remove,
    size,
  }
}


const topKFrequent = function (nums = [], k) {
  const map = nums.reduce((acc, num) => {
    if (acc[num]) {
      acc[num] += 1
    } else {
      acc[num] = 1
    }
    return acc
  }, {})
  const maxHeap = Heap((a, b) => map[b] - map[a])(
    Object.keys(map).map(x => parseInt(x)),
  )
  const res = []
  for (let i = 1; i <= k; i++) {
    res.push(maxHeap.extract())
  }
  return res
}