[LeetCode] Top K Frequent Elements 개인풀이

문제

숫자 배열 nums과 숫자 k가 주어집니다. nums 배열에서 자주 반복되는 값을 k개 반환하도록 합니다. 순서는 상관 없습니다.

예제_1

입력: nums = [1,1,1,2,2,3], k = 2
출력: [1,2]

예제_2

입력: nums = [1], k = 1
출력: [1]

제약조건

• 1 <= nums.length <= 10^5
• k 의 범위 1 ~ nums에 있는 유니크한 숫자의 수
• 정답은 유니크하다고 보장

나아가기

배열의 사이즈가 N 일때 알고리즘의 시간복잡도는 O(NlogN) 보다 적은 시간이 걸리자.

문제 원본 : https://leetcode.com/problems/top-k-frequent-elements/

개인 풀이

저는 이 문제를 입력의 값을 ( element 값, counting )의 형태로 묶어낸다음 그것을 counting을 기준으로 내림차순 정렬을 하여 앞에서부터 k번째까지의 숫자까지 가져오는 알고리즘을 생각하였습니다.

저는 이것을 TreeSet에서 key값이 겹치면 그 값이 대체된다는 점을 이용하여 해결하려고 했습니다. 하지만 TreeSet에 들어있는 값을 가져와 업데이트 하는 형식을 생각했지만... TreeSet은 내부의 값을 가져올 수 없다는 것을 알게 되어 TreeMap으로 대체하여 구현하였습니다. 제가 구현한 로직은 아래와 같습니다.

1. TreeMap에 이용할 Node를 만들며 key 기준을 value로 잡는다.
2. nums를 순회하며 treeMap에 넣으며 값이 이미 존재하면 count를 1 올리며, 없으면 추가한다.
3. keySet을 List로 가져온 후 count를 기준으로 정렬한다.
4. k의 갯수만큼 뒤에서 가져와 리턴한다.

코드

class Solution {
    public int[] topKFrequent(int[] nums, int k) {

        int[] result = new int[k];
        TreeMap<Node, Node> treeMap = new TreeMap<>();

        for (int index = 0; index < nums.length; index++) {
            Node nodeKey = new Node(nums[index]);
            Node nodeValue = treeMap.get(nodeKey);

            if (nodeValue != null) { // 값이 있으면 1 추가후 nodeValue로 교체
                nodeValue.count++;
                treeMap.put(nodeValue, nodeValue);
            } else {
                treeMap.put(nodeKey, nodeKey); // 값이 없으면 새로만든 nodeKey 추가
            }
        }

        ArrayList<Node> semiResult = new ArrayList<>(treeMap.keySet()); // TreeMap의 Key를 List로 뽑아내기
        semiResult.sort(Comparator.comparing(node -> node.count)); // count로 정렬

        for(int index = 0; index < result.length; index++) {
            result[index] = semiResult.get(semiResult.size() - index - 1).value;
        }

        return result;

    }
}

class Node implements Comparable<Node> {

    int value;

    int count;

    public Node(int value) {
        this.value = value;
        this.count = 1;
    }

    @Override
    public int compareTo(Node o) {
        return Integer.compare(value, o.value); // 비교대상 Key 값 선정
    }
}

LeetCode 솔루션

개인적으로 문제를 풀었으나 코드가 세련되지 못하다고 생각했습니다. 문제에 대한 동일한 접근방법을 가지고 좀 더 깔끔하게 코딩된 것이 Solution으로 제공되고 있어 공유드립니다. 우선순위큐를 만들때 비교방법을 외부의 파라미터를 이용할 수 있다는 사실을 알 게 되었습니다.

1. HashMap에 값을 (key = value, value = count)의 형태로 넣습니다.
2. 우선순위 큐를 만들어냅니다. 우선순위큐를 만들때 비교하는 방법은 hashMap의 value를 가져와서 비교하는 것입니다.
3. 우선 순위큐에서 k만큼 값을 뽑아서 리턴합니다.

코드

public int[] topKFrequent_1(int[] nums, int k) {

    // O(1) time
    if (k == nums.length) {
        return nums;
    }

    // 1. build hash map : character and how often it appears
    // O(N) time
    Map<Integer, Integer> count = new HashMap();
    for (int n : nums) {
        count.put(n, count.getOrDefault(n, 0) + 1);
    }

    // init heap 'the less frequent element first'
    // 이렇게 하면 heap.add(n)만 하더라도 hashMap의 value를 통해 비교할 수 있음.
    Queue<Integer> heap = new PriorityQueue<>((n1, n2) -> count.get(n1) - count.get(n2));

    // 2. keep k top frequent elements in the heap
    // O(N log k) < O(N log N) time
    for (int n : count.keySet()) {
        heap.add(n);
        if (heap.size() > k) heap.poll();
    }

    // 3. build an output array
    // O(k log k) time
    int[] top = new int[k];
    for (int i = k - 1; i >= 0; --i) {
        top[i] = heap.poll();
    }
    return top;
}

마무리

오늘은 이렇게 Top K Frequent elements의 문제풀이를 해보았습니다.

우선순위큐를 사용할 수 있는 새로운 방법을 배울 수 있어서 좋았습니다.

감사합니다.