347. Top K Frequent Elements

package LeetCode_347

import java.util.*
import kotlin.collections.ArrayList
import kotlin.collections.HashMap

/**
 * 347. Top K Frequent Elements
 * https://leetcode.com/problems/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]

Time complexity try on O(n)
 * */
class Solution {
    fun topKFrequent(nums: IntArray, k: Int): List<Int> {
        val result = ArrayList<Int>()
        //first: num
        //second: the frequency of num
        val maxHeap = PriorityQueue<Pair<Int,Int>>(kotlin.Comparator { a, b -> b.second - a.second })
        val frequencyMap = HashMap<Int, Int>()
        for (num in nums) {
            frequencyMap.put(num, frequencyMap.getOrDefault(num, 0) + 1)
        }
        frequencyMap.forEach { key, value ->
            //value is frequency
            val pair = Pair(key,value)
            maxHeap.add(pair)
        }
        for (i in 0 until k){
            result.add(maxHeap.remove().first)
        }
        return result
    }
}