백준 14427 - 수열과 쿼리 15

문제

https://www.acmicpc.net/problem/14427

리뷰

TreeMap과 HashMap을 이용하여 풀이할 수 있는 문제였다.

우선 수열의 원소를 표현하기 위해 인덱스와 값을 필드로 가지는 Element 클래스를
정의하였다. 또한, $N$개의 원소중에 가장 작은 값(값이 같을 때는 가장 작은 인덱스)를
가지는 원소를 $O(logN)$의 복잡도내에 조회하기 위해 TreeMap<Element, Integer>를
활용하여 키로는 Element, 값으로는 인덱스를 저장했다.

한편, 1쿼리가 주어졌을 때 TreeMap에서 인덱스를 기반으로 해당 원소를
갱신해주어야 한다. 하지만, TreeMap 의 키가 Element 이기 때문에 탐색을
위해 추가적으로 다른 자료 구조를 활용하거나 하는 방법이 필요했다. 이에 따라
Element에 equals()와 hashCode()를 오버라이딩하고
HashMap<Integer, Element>를 통해 인덱스를 기반으로 Element 데이터를
$O(1)$의 시간내에 조회할 수 있도록 구성했다.

이런 설계를 바탕으로 1쿼리가 주어졌을 때

• HashMap에서 인덱스를 기반으로 Element를 조회
• TreeMap에서 해당 원소를 삭제한 후
• TreeMap과 HashMap에 새로운 원소를 투입

위 절차에 따라 연산을 $O(logN)$ 복잡도 안에 처리할 수 있었다.

로직의 시간복잡도는 쿼리를 처리하는 부분에서 $O(MlogN)$ 복잡도로 수렴하므로
$M=100,000$, $N=100,000$인 최악의 경우에도 무난히 제한 조건 1초를 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Objects;
import java.util.StringTokenizer;
import java.util.TreeMap;

import static java.lang.Integer.*;

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = parseInt(br.readLine());
        TreeMap<Element, Integer> map = new TreeMap<>((e1, e2) -> {
            if (e1.value == e2.value) return compare(e1.idx, e2.idx);

            return compare(e1.value, e2.value);
        });
        HashMap<Integer, Element> hashMap = new HashMap<>();

        StringTokenizer st = new StringTokenizer(br.readLine());
        int value;
        Element e;
        for (int idx = 1; idx <= N; idx++) {
            value = parseInt(st.nextToken());
            e = new Element(idx, value);
            map.put(e, value);
            hashMap.put(idx, e);
        }

        int M = parseInt(br.readLine());
        int cmd, i, v;
        StringBuilder sb = new StringBuilder();
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            cmd = parseInt(st.nextToken());

            if (cmd == 1) {
                i = parseInt(st.nextToken());
                v = parseInt(st.nextToken());

                e = hashMap.get(i);
                map.remove(e);
                e = new Element(i, v);
                map.put(e, v);
                hashMap.put(i, e);
            } else {
                sb.append(map.firstKey().idx).append("\n");
            }
        }

        System.out.print(sb);
        br.close();
    }


    static class Element {
        int idx, value;

        public Element(int idx, int value) {
            this.idx = idx;
            this.value = value;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Element element = (Element) o;
            return idx == element.idx && value == element.value;
        }

        @Override
        public int hashCode() {
            return Objects.hash(idx, value);
        }
    }
}

결과