use rand::Rng;
use std::{collections::HashMap, iter::repeat};
impl Solution {
    pub fn top_k_frequent(nums: Vec<i32>, k: i32) -> Vec<i32> {
        let mut map: HashMap<i32, i32> = HashMap::new();
        for num in nums.iter() {
            map.insert(
                *num,
                if map.contains_key(num) {
                    *map.get(num).unwrap()
                } else {
                    0
                } + 1,
            );
        }
        let mut i = map.len() as i32;
        let mut arr: Vec<Vec<i32>> = repeat(repeat(0).take(2).collect())
            .take(i as usize)
            .collect();
        for (key, value) in map.iter() {
            i -= 1;
            arr[i as usize][0] = *key;
            arr[i as usize][1] = *value;
        }
        let arr_len = arr.len() as i32;
        Solution::more_less(&mut arr, 0, arr_len - 1, k);
        let mut ans: Vec<i32> = repeat(0).take(k as usize).collect();
        while i < k {
            ans[i as usize] = arr[i as usize][0];
            i += 1;
        }
        return ans;
    }

    fn more_less(arr: &mut Vec<Vec<i32>>, l: i32, r: i32, k: i32) {
        if k == r - l + 1 {
            return;
        }
        arr.swap(
            r as usize,
            (l + rand::thread_rng().gen_range(0, r - l + 1)) as usize,
        );
        let pivot = Solution::partition(arr, l, r);
        if pivot - l == k {
            return;
        } else if pivot - l > k {
            Solution::more_less(arr, l, pivot - 1, k);
        } else {
            Solution::more_less(arr, pivot, r, k - pivot + l);
        }
    }

    fn partition(arr: &mut Vec<Vec<i32>>, l: i32, r: i32) -> i32 {
        let mut left = l - 1;
        let mut index = l;
        while index < r {
            if arr[index as usize][1] <= arr[r as usize][1] {
                index += 1;
            } else {
                left += 1;
                arr.swap(left as usize, index as usize);
                index += 1;
            }
        }
        left += 1;
        arr.swap(left as usize, r as usize);
        return left;
    }
}

fn main() {
    let num2 = vec![1, 1, 1, 2, 2, 3];
    let k = 2;
    let ans = Solution::top_k_frequent(num2, k);
    println!("ans = {:?}", ans);
}

struct Solution {}