LeetCode-in-Rust

347. Top K Frequent Elements

Medium

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Example 1:

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

Output: [1,2]

Example 2:

Input: nums = [1], k = 1

Output: [1]

Constraints:

Follow up: Your algorithm’s time complexity must be better than O(n log n), where n is the array’s size.

Solution

use std::collections::HashMap;

impl Solution {
    pub fn top_k_frequent(nums: Vec<i32>, k: i32) -> Vec<i32> {
        nums.iter().fold(HashMap::new(), |mut map, n| {
            let mut counter = map.entry(n).or_insert(0);
            *counter += 1;
            map
        })
        .drain()
        .fold(vec![(0, 0); k as usize], |mut top_k, (&num, count)| {
            if count > top_k[0].1 {
                top_k[0] = (num, count);

                let mut next_index = 1;
                while next_index < k as usize {
                    if count > top_k[next_index].1 {
                        let temp = top_k[next_index];
                        top_k[next_index] = (num, count);
                        top_k[next_index - 1] = temp;
                    }
                    next_index += 1;
                }
            }

            top_k
        })
        .into_iter()
        .map(|(num, count)| num)
        .collect()
    }
}