[Boj_1194] 아이스크림 도둑 지호

📎 문제 링크

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

 

▸ 문제

지호는 매일 아이스크림 가게에 방문한다. 아이스크림을 먹던 지호는 놀라 자빠질 수밖에 없었다. 실수로 민트초코 맛을 먹었기 때문이다. 대다수의 사람은 치약 맛이 난다는 이유로 민트초코를 싫어한다. 아이스크림으로 이를 닦는다는 발상은 누가 한 것인지 궁금할 뿐이다. 아무튼 매번 아이스크림을 사 먹는 것이 지겨워진 지호는 이제부터 아이스크림을 훔쳐 먹기로 결심하였다.

아이스크림 가게에는 다양한 맛의 아이스크림 N개가 한 줄로 배치되어 있다. 아이스크림에는 번호가 매겨져 있는데, 가장 왼쪽 아이스크림이 1번, 그 오른쪽은 2번, ..., 가장 오른쪽 아이스크림은 N번이다. 지호는 항상 양이 가장 많은 아이스크림을 선택하여 전부 먹는다. 양이 가장 많은 아이스크림이 여러 개라면 가장 왼쪽에 있는 것을 먹는다.

지호는 대다수의 사람과 마찬가지로 민트초코 맛을 싫어한다. 다행히 지호는 아이스크림의 양이 주어질 때 아이스크림의 맛을 알 수 있다. 지호의 판별법에 따르면, 아이스크림의 양이 7의 배수라면 민트초코 맛이고, 그렇지 않다면 민트초코 맛이 아니라고 한다.

지호는 민트초코를 싫어한다는 사실을 명심하라. 민트초코 맛 아이스크림을 먹은 지호는 크게 분노하여 남아 있는 아이스크림의 순서를 좌우로 뒤집는다. 즉, K개의 아이스크림이 있다면 i번째 아이스크림과 (K - i + 1)번째 아이스크림의 위치를 뒤바꾼다. (1 ≤ i ≤ ⌊K / 2⌋)

지호는 N개의 아이스크림 중 M개의 아이스크림을 먹으려 한다. 아이스크림의 양이 주어졌을 때, 지호가 먹은 아이스크림의 번호를 구하는 프로그램을 작성하시오.

 

▸ 입력

첫 번째 줄에 전체 아이스크림의 개수 N과 지호가 먹을 아이스크림의 개수 M이 주어진다.

두 번째 줄에 N개의 정수 A1, A2, ..., AN이 주어진다. 이때 Ai는 i번 아이스크림의 양을 의미한다.

모든 입력은 공백으로 구분되어 주어진다.

 

▸ 출력

M개의 줄에 걸쳐 i(1 ≤ i ≤ M)번째 줄에는 지호가 i번째로 먹은 아이스크림의 번호를 출력한다.

 

▸ 제한

• 1 ≤ N ≤ 100,000
• 1 ≤ M ≤ N
• 1 ≤ Ai ≤ 1,000,000,000

 

📍 문제 정보

• 🥇  문제 레벨 : 골드 4
• 🔔  문제 유형 : 자료 구조, 정렬, 덱
• 💬  풀이 언어 : JAVA

 

🤔 문제 풀이

이 문제는 TreeMap과 TreeSet을 활용해서 해결했다.

 

먼저, TreeMap을 선언한 뒤 key, value는 각각 다음과 같이 설정한다.

• key : 아이스크림 양 (amount)
• value : 그 양을 가진 아이스크림들의 인덱스들 (TreeSet<Integer>)

이때, TreeSet<Integer>을 사용하는 이유는 ? 🤔

 

문제 조건에 따르면 '양이 가장 많은 아이스크림이 여러 개라면 가장 왼쪽에 있는 것을 먹는다.'

따라서 같은 양의 아이스크림이 여러 개 있을 경우,

가장 작은 인덱스(가장 왼쪽)를 가진 아이스크림을 먹어야 한다.

 

TreeSet<Integer>은 인덱스를 자동으로 오름차순 정렬해서 저장하므로,

• pollFirst() → 가장 왼쪽 아이스크림
• pollLast() → reverse 상태일 때 가장 오른쪽 아이스크림
• 을 각각 빠르게 꺼낼 수 있다.

 

💡 참고로 코드에서 사용한 putIfAbsent()  메서드에 대해 알아보자면, 

putIfAbsent(key, value)는 Java 8 이상에서 지원되며,

 

• 해당 키가 존재하지 않으면 value를 삽입하고
• 이미 존재하면 아무것도 안 함

예를 들어, 어떤 양이 14인 아이스크림이 처음 등장했다고 가정하자.

if (!iceCreams.containsKey(14)) {
    iceCreams.put(14, new TreeSet<>());
}
iceCreams.get(14).add(i);

이 코드를 간단히 바꿀 수 있다.

iceCreams.putIfAbsent(14, new TreeSet<>());
iceCreams.get(14).add(i);

즉, amount 양의 아이스크림 그룹(TreeSet<index>)이 없으면 새로 만들고, 있으면 기존 걸 사용해서 해당 인덱스 i를 추가하는 구조이다.

 

위 풀이의 코드를 바탕으로 전체적인 코드를 작성해 보면 아래와 같다.

 

✔ 코드

메모리 : 65092KB
시간 : 620ms

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
import java.util.TreeMap;
import java.util.TreeSet;

public class Main {
    static boolean isReverseOrder = false;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken()); // 전체 아이스크림의 개수
        int m = Integer.parseInt(st.nextToken()); // 지호가 먹을 아이스크림의 개수

        // 아이스크림의 양, 그 양을 가진 아이스크림의 인덱스들
        TreeMap<Integer, TreeSet<Integer>> iceCreams = new TreeMap<>();
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            int amount = Integer.parseInt(st.nextToken());
            // 해당 키가 존재하지 않으면 value를 삽입하고
            // 이미 존재하면 아무것도 안 함
            iceCreams.putIfAbsent(amount, new TreeSet<>());
            iceCreams.get(amount).add(i);
        }

        StringBuilder sb = new StringBuilder();
        while (m --> 0) {
            TreeSet<Integer> group = iceCreams.lastEntry().getValue(); // 가장 양 많은 그룹

            if (isReverseOrder) { // 좌우로 뒤집었다면
                sb.append(group.pollLast()).append("\n"); // 오른쪽에서 먹기
            } else {
                sb.append(group.pollFirst()).append("\n"); // 왼쪽에서 먹기
            }

            if (iceCreams.lastKey() % 7 == 0) { // 민초를 먹었다면
                isReverseOrder = !isReverseOrder;
            }

            if (group.isEmpty()) {
                iceCreams.pollLastEntry(); // 먹은 그룹 다 비우면 삭제
            }
        }

        System.out.println(sb.toString());

        br.close();
    }
}

 

처음에 Deque를 사용해서 문제를 풀었는데 계속 시간 초과가 떴다..

도저히 아이디어가 안 떠올라서 다른 분의 풀이를 참고해서 풀었다. 오랜만에 TreeMap과 TreeSet을 사용해 문제를 풀었는데,

처음 보는 메서드에 대해서도 공부할 수 있었다. 어쩌면 정렬 문제가 제일 어려울지도 모르겠다...⭐