[백준 -32510번] 키가 비슷한 친구 - Java

문제 링크

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

문제

티어: 골드 2
알고리즘: 정렬, 우선순위 큐

$N$ 명의 사람들이 한 줄로 서 있다. 각 사람은 왼쪽에서 오른쪽으로 순서대로 $1, 2, \ldots, N$ 의 번호가 붙어 있다. $i$ 번 사람의 키는 $A_i$ 이다. $i$ 번 사람과 $j$ 번 사람의 거리는 $|i - j|$ 이다.

모든 $i$ 에 대해서, 키가 $A_i - K$ 이상 $A_i$ 이하면서, 나의 오른쪽에 서 있지 않은 사람들 중 가장 먼 사람을 찾고, 그 사람과의 거리를 합한 것을 출력하라. $i$ 번 사람에 대해 자신은 위 조건을 만족하기 때문에, 항상 그러한 사람들은 존재한다.

입력

파일의 첫째 줄에 테스트 케이스의 개수를 나타내는 자연수 $T$ 가 주어지고,

이후 차례로 $T$ 개의 테스트 케이스가 주어진다. ($1 \le T \le 31$)

각 테스트 케이스의 첫 줄에는 정수 $N, K$ 가 주어진다. ($1 \le N \le 200\,000, 0 \le K \le 500\,000$)

다음 줄에는 $N$ 개의 정수 $A_1, \ldots, A_N$ 이 주어진다. ($0 \le A_i \le 500\,000$)

모든 테스트 케이스들에 대한 $N$ 의 합은 $2\,000\,000$ 이하이다.

출력

각 테스트 케이스마다 첫 줄에는 Case #$C$ 를 출력하여야 한다. 이때 $C$는 테스트 케이스의 번호이다.

다음 줄에는 문제의 정답을 출력한다.

풀이

모든 $i$ 에 대해서, 키가 $A_i - K$ 이상 $A_i$ 이하면서, 나의 오른쪽에 서 있지 않은 사람들 중 가장 먼 사람을 찾고, 그 사람과의 거리를 합한 것을 출력해야 한다.

현재 i에 대해서 범위를 구하고, 내 왼쪽에서 그 범위안에 들어가는 것 중 가장 먼 j를 찾아야 하는데, 단순히 선형 탐색으로 하면 시간 초과가 난다.

이는 매번 범위가 바뀌기 때문에 가능한 왼쪽의 값이 불규칙적으로 빠지고, 들어가고를 반복하기 때문이다.

그러면 범위를 고정하는 것은 불가능하지만, 규칙적인 방법이 있지 않을까? 생각을 해봐야 한다.

사람들을 A를 기준으로 오름차순 해준다면??

5 1 3 5 8 6 6 9 10는
1 3 5 5 6 6 8 9 10이 되는데,

변화는 범위를 보자, (K는 3)

1은 1 - 3 <= x <= 1
3는 3 - 3 <= x <= 3
5는 5 - 3 <= x <= 5
...

1의 범위로 포함했을 때 1 하나만 있고,
3의 범위로 포함했을 때 1 3 이렇게 둘이 있고,
5의 범위로 포함했을 때는 3, 5가 된다. 2부터 시작하기 때문에 1이 빠진 것을 볼 수 있다.

여기서 전의 poll을 유지하면서, min보다 작은 부분은 빼주면, 다시는 빼준 값이 들어갈 일이 없다. 오름차순으로 범위가 조정되기 때문이다.

그러면 중요한 것은 A를 기준으로 정렬했기 때문에 순서는 뒤죽박죽일 것이다. 이때 왼쪽에 가장 먼 사람을 구하려면 우선순위 큐를 활용할 수 있다.

index가 작은 순으로 정렬되게끔 우선순위 큐를 활용한다면, pq에 peek()가 최초로 조건을 만족했을 때 그 값은 조건을 만족하며 가장 왼쪽의 값임을 보장할 수 있다.

매번 자기자신을 pq에 넣기 때문에 자기자신보다 왼쪽에 있는 조건을 만족하는 값이 없다면, 결국 자기자신이 조건을 만족하며 가장 왼쪽에 있는 값이 될 것이다.

중요한건 오름차순 순으로 사람을 탐색하기 때문에 pq.peek()가 min (Ai - K)보다 작다면 과감하게 빼야한다는 것이다. 앞으로도 그 값은 사용될 일이 없기 때문이다.

또한 만약 A값이 같다면 index가 빠른 순으로 정렬되어야 한다. 1 3 5 5의 경우를 생각해 보면 된다.

자세한 구현 사항은 코드를 확인하길 바란다.

이 풀이의 시간 복잡도는 O(NlogN)이다.

소스 코드

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

class Person implements Comparable<Person> {
    int A, ind;
    Person(int ind, int A) {
        this.ind = ind;
        this.A = A;
    }
    
    @Override
    public int compareTo(Person o) {
        if(this.ind < o.ind) {
            return -1;
        } else if(this.ind > o.ind) {
            return 1;
        }
        return 0;
    }
}

public class Main {
    static int T;
  public static void main(String args[]) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      T = Integer.parseInt(br.readLine());
      
      StringBuilder sb = new StringBuilder();
      for(int t=1; t<=T; t++) {
          StringTokenizer st = new StringTokenizer(br.readLine());
          int N = Integer.parseInt(st.nextToken());
          int K = Integer.parseInt(st.nextToken());
          
          StringTokenizer st2 = new StringTokenizer(br.readLine());
          ArrayList<Person> list = new ArrayList<>();
          for(int i=1; i<=N; i++) {
              list.add(new Person(i, Integer.parseInt(st2.nextToken())));
          }
          
          Collections.sort(list, new Comparator<Person>() {
              @Override
              public int compare(Person p1, Person p2) {
                  if(p1.A < p2.A) {
                      return -1;
                  } else if(p1.A == p2.A) {
                      if(p1.ind < p2.ind) {
                          return -1;
                      } else if(p1.ind > p2.ind) {
                          return 1;
                      }
                  } else {
                      return 1;
                  }
                  return 0;
              }
          });
          
          PriorityQueue<Person> pq = new PriorityQueue<>();
          
          long answer = 0;
          for(int i=0; i<list.size(); i++) {
              pq.add(list.get(i)); //먼저 현재 값을 넣어줌.
              int min = list.get(i).A - K; //최소 조건 이 값보단 크거나 같아야 됨.
              
              while(!pq.isEmpty()) {
                  //pq poll을 일단 조건에 맞게 조정
                  if(pq.peek().A >= min) {
                      answer += list.get(i).ind - pq.peek().ind;
                      break;
                  }
                  pq.poll();
              }
          }
          sb.append("Case #" + t).append("\n");
          sb.append(answer).append("\n");
      }
      System.out.println(sb.toString().trim());
  }
}