샘터 18513

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

너비탐색이라는 힌트를 얻어서 풀 수 있었다

O(n log n)

우선 순위 큐를 사용하여 시간 복잡도가 log n인 작업을 수행하고 있습니다. 우선 순위 큐는 힙(heap)으로 구현되어 있어서 삽입과 삭제 연산이 log n의 시간 복잡도를 가집니다.
n개의 요소에 대해 우선 순위 큐에 삽입하는 작업을 하므로, 전체적으로 n log n의 시간 복잡도를 가집니다.
HashSet 연산은 시간 복잡도가 O(1)이므로 큰 영향을 미치지 않습니다.

import java.io.*;
import java.util.*;

class Zip implements Comparable<Zip>{
    int pos;
    int dist;
    Zip(int pos, int dist){
        this.pos = pos;
        this.dist = dist;
    }

    @Override
    public int compareTo(Zip o){
        return this.dist < o.dist ? -1 : 1;
    }
}
public class Main {
    public static void main(String[] args) throws IOException{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        String[] input = br.readLine().split(" ");

        int n = Integer.parseInt(input[0]);
        int k = Integer.parseInt(input[1]);

        input = br.readLine().split(" ");

        PriorityQueue<Zip> pq = new PriorityQueue<>();
        HashSet<Integer> visit = new HashSet<>();

        for (int i = 0; i < n; i++) {
            int temp = Integer.parseInt(input[i]);
            pq.add(new Zip(temp, 0));
            visit.add(temp);
        }

        long answer = 0;
        int zipcnt = k;
        while(zipcnt > 0){

            Zip cur = pq.poll();

            int neighbor = 0;
            for (int i = 0; i < 2; i++) {

                if(i == 0)
                    neighbor = cur.pos-1;
                else
                    neighbor = cur.pos+1;

                if(visit.contains(neighbor))
                    continue;

                visit.add(neighbor);
                zipcnt--;
                answer += cur.dist+1;
                pq.add(new Zip(neighbor, cur.dist+1));

                if(zipcnt <= 0)
                    break;
            }
        }

        System.out.println(answer);

    }

}