use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};

pub struct Solution;

impl Solution {
    pub fn top_k_frequent(words: Vec<String>, k: i32) -> Vec<String> {
        let mut words_map = HashMap::new();
        let mut heap = BinaryHeap::new();
        for word in words {
            *words_map.entry(word).or_insert(0) += 1;
        }
        for (idx, entry) in words_map.into_iter().enumerate() {
            let word_count = WordCount::new(entry.0, entry.1);
            heap.push(word_count);
            if idx >= k as usize {
                heap.pop();
            }
        }
        heap.into_sorted_vec()
            .into_iter()
            .map(|word_count| word_count.word)
            .collect()
    }
}

struct WordCount {
    word: String,
    count: i32,
}

impl WordCount {
    fn new(word: String, count: i32) -> Self {
        WordCount { word, count }
    }
}

impl PartialEq for WordCount {
    fn eq(&self, other: &Self) -> bool {
        self.count.eq(&other.count) && self.word.eq(&other.word)
    }
}

impl Eq for WordCount {}

impl PartialOrd for WordCount {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(match other.count.cmp(&self.count) {
            Ordering::Equal => self.word.cmp(&other.word),
            order => order,
        })
    }
}

impl Ord for WordCount {
    fn cmp(&self, other: &Self) -> Ordering {
        match other.count.cmp(&self.count) {
            Ordering::Equal => self.word.cmp(&other.word),
            order => order,
        }
    }
}