leetcode 432. All O`one Data Structure

과거 지원했던 한 회사의 라이브 코딩 테스트에서 나왔던 문제였는데, 제대로 대답하지 못했던 기억이 난다. 다시 풀고 잊지 않기 위해 기록해 둔다.

 

이 문제의 관건은 결국 최대값과 최소값을 O(1)에 접근하는 것이다. 내가 알고 있는 최대값과 최소값을 가장 빠르게 찾을 수 있는 것은 heap 을 쓰는 것인데 heap을 쓴다고 하더라도 O(logN)이 걸리기 때문에 이 문제를 풀 수 없다.

 

문제를 푸는 실마리는 count 값이 1씩 증감한다는 것이다. 따라서 count 값이 변경될 때 순서를 알기 위해서 전체 노드를 다 살펴볼 필요 없이 자기자신보다 1 크거나 작은 노드만 살펴보면 된다. 자연스럽게 linked list 를 생각해 볼 수 있다. linked list 는 노드의 추가 삭제가 상수시간에 이뤄지고, 다음 노드를 가르키는 포인트를 가지고 있으므로 이 문제를 해결하는데 적합한 자료 구조라고 할 수 있다.

 

count 가 증가할수도 감소할 수도 있기 때문에 각 노드는 이전과 다음 노드의 포인트를 모두 가지고 있어야 비교가 가능하다. 이 때 사요하는것이 double linked list 이다. 다음은 hashmap 과 double linked list 를 사용해서 문제를 푼 코드이다. 언어는 java 를 사용했다.

import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

public class AllOne {

    // 각 Node 를 담는 클래스
    private class Node {
        int count; // 노드에서 나타내는 문자열 빈도수
        Set<String> keys; // 같은 count 값을 가진 문자열 조합
        Node prev; // 이전 노드
        Node next; // 다음 노드

        Node(int count) {
            this.count = count; // count 값을 받아서 새로운 노드 생성함
            this.keys = new LinkedHashSet<>(); // 삽입 순서를 유지하는 집합으로 만듦(count 값이 같은 경우 때문)
        }
    }

    private Map<String, Node> keyToNode; // 각 문자열이 속한 Node 를 매핑
    private Node head; // 시작 노드(더미)
    private Node tail; // 마지막 노드(더미)

    public AllOne() {
        this.keyToNode = new HashMap<>();
        this.head = new Node(Integer.MIN_VALUE); // 시작 더미노드, count 는 최소값을 가짐
        this.tail = new Node(Integer.MAX_VALUE); // 마지막 더미 노드, count 는 최대값을 가짐
        this.head.next = this.tail;
        this.tail.prev = this.head;
    }

    // 문자열 count 를 1 증가
    public void inc(String key) {
        // 이미 존재하는 key 값이 경우
        if (keyToNode.containsKey(key)) {
            Node node = keyToNode.get(key); // 매핑된 노드를 가져온다.
            moveKeyAfter(key, node, node.count + 1); // 새로운 카운트는 1 증가
            return;
        }
        // key 가 존재하지 않는 경우
        if (head.next.count != 1) {
            // head 의 다음 노드가 1 이 아니면 새로운 노드를 생성
            addNodeAfter(new Node(1), head);
        }
        updateKey(key, head.next);
    }

    private void moveKeyAfter(String key, Node node, int newCount) {
        Node targetNode; // 새로운 카운트에 매핑된 key 값을 담을 targetNode
        if (node.next.count != newCount) {
            // 새로운 카운트 값을 가진 노드가 존재하지 않으면
            targetNode = new Node(newCount); // 새로운 노드 생성
            addNodeAfter(targetNode, node); // 새로운 노드를 기존 노드의 뒤에 추가
        } else {
            // 새로운 카운트 값을 가진 노드가 이미 존재하는 경우
            targetNode = node.next; // 기존 노드의 다음 노드가 새 노드가 된다.
        }
        updateKey(key, targetNode); // 새로운 노드 정보 업데이트
        removeKeyFromNode(key, node); // 기존 노드에서 key 제거
    }

    private void moveKeyBefore(String key, Node node, int newCount) {
        Node targetNode; // 새로운 카운트에 매핑된 key 값을 담을 targetNode
        if (node.prev.count != newCount) {
            // 새로운 카운트 값을 가진 노드가 존재하지 않으면
            targetNode = new Node(newCount); // 새로운 노드 생성
            addNodeBefore(targetNode, node); // 새로운 노드를 기존 노드의 앞에 추가
        } else {
            // 새로운 카운트 값을 가진 노드가 이미 존재하는 경우
            targetNode = node.prev; // 기존 노드의 이전 노드가 새 노드가 된다.
        }
        updateKey(key, targetNode); // 새로운 노드 정보 업데이트
        removeKeyFromNode(key, node); // 기존 노드에서 key 제거
    }

    // 새로운 노드를 기존 노드의 앞에 추가
    private void addNodeBefore(Node targetNode, Node node) {
        targetNode.next = node;
        targetNode.prev = node.prev;
        node.prev.next = targetNode;
        node.prev = targetNode;
    }

    // node 에 존재하는 key 값 제거
    private void removeKeyFromNode(String key, Node node) {
        node.keys.remove(key); // 현재 count 값에 매핑된 key 값을 제거
        if (node.keys.isEmpty()) {
            // node에 key가 하나도 없으면 노드를 제거
            removeNode(node);
        }
    }

    private void removeNode(Node node) {
        node.prev.next = node.next;
        node.next.prev = node.prev;
    }

    private void updateKey(String key, Node targetNode) {
        targetNode.keys.add(key); // 새 노드에 key 를 추가한다.
        keyToNode.put(key, targetNode); // map 을 새로운 노드로 업데이트
    }

    // 새로운 노드를 기존 노드의 뒤에 추가
    private void addNodeAfter(Node targetNode, Node node) {
        targetNode.prev = node;
        targetNode.next = node.next;
        node.next.prev = targetNode;
        node.next = targetNode;
    }

    // 문자열 카운트를 1 감소
    public void dec(String key) {
        if (!keyToNode.containsKey(key)) {
            // key 값에 매핑된 노드가 없으면 아무것도 하지 않음
            return;
        }
        Node node = keyToNode.get(key);
        if (node.count == 1) {
            // 만약 현재 카운트가 1 이면 노드를 제거
            keyToNode.remove(key);
            removeKeyFromNode(key, node);
            return;
        }
        // count 가 1 보다 크면
        moveKeyBefore(key, node, node.count - 1); // 카운트 1 감소
    }

    // count 가 최대인 값 반환
    public String getMaxKey() {
        return tail.prev == head ? "" : tail.prev.keys.iterator().next();
        // tail.prev가 head면 데이터가 없으므로 빈 문자열 반환
        // 그렇지 않으면 tail.prev의 keys 중 하나 반환
    }

    // 최소 count를 가진 문자열 중 하나를 반환
    public String getMinKey() {
        return head.next == tail ? "" : head.next.keys.iterator().next();
        // head.next가 tail이면 데이터가 없으므로 빈 문자열 반환
        // 그렇지 않으면 head.next의 keys 중 하나 반환
    }
}