[JAVA] 백준 3078 좋은 친구

 

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

 

문제분석

이 풀이는 “이름 길이가 같은 학생들끼리만 좋은 친구 후보가 된다”는 점을 그대로 자료구조에 반영한 방식이다. 이름 길이는 2~20이므로 길이별로 총 19개의 큐를 두고(ArrayDeque<Integer>[19]), 각 큐에는 그 길이를 가진 학생들의 등수(인덱스) 만 저장한다. 학생을 성적순으로 0번부터 읽어가면서, 현재 학생의 이름 길이에 해당하는 큐만 확인하면 된다. 즉, 서로 다른 길이끼리는 애초에 비교 대상이 아니니 큐도 분리해두는 게 핵심이다.

 

각 학생 i에 대해 같은 길이 큐를 보면서, 큐의 맨 앞(가장 오래된 등수)부터 i - peek > K 인 애들은 전부 탈락이므로 poll()로 제거한다. 이렇게 하면 큐에는 항상 현재 학생 기준으로 K등수 이내에 있는 같은 길이 학생들만 남는다. 그 상태에서 result += 큐의 크기 를 해주면, 지금 학생 i와 좋은 친구가 될 수 있는 기존 학생 수를 한 번에 더한 셈이 된다. 마지막으로 현재 학생의 등수 i를 offer()로 넣어 다음 학생들의 비교 대상에 포함시키면 끝이다. 결국 각 학생은 큐에 한 번 들어가고 한 번 빠지므로 전체가 선형에 가깝게 돌아가고, 길이별로 윈도우를 유지하는 슬라이딩/큐 운영으로 깔끔하게 해결한 형태다.

코드

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

public class Main {
    static int K; // 등수 차이, K 등수 차이를 넘으면 좋은 친구 탈락
    static long result = 0; // 좋은 친구 수

    // 학생 글자수별 queue 운영(2글자 ~ 20글자 이므로, 총 19개면 됨). idx는 0 은 2글자, 1은 3글자..
    static ArrayDeque<Integer>[] q; // 각 q별 저장하는 값은 "해당 친구 등수"

    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()); // 반 학생의 이름(성적순) 수
        K = Integer.parseInt(st.nextToken());

        q = new ArrayDeque[19];
        for (int i = 0; i < 19; i++) {
            q[i] = new ArrayDeque<>();
        }

        for (int i = 0; i < N; i++) {
            String s = br.readLine();
            int len = s.length();
            count(i, len);
        }
        System.out.println(result);
    }

    static void count(int prizeNum, int len) {
        int qIdx = len - 2;// q는 0부터 시작하기 때문에, 2글자 이름이면 0
        if (q[qIdx].size() == 0) {
            q[qIdx].offer(prizeNum);
        } else {
            while (!q[qIdx].isEmpty()) {
                if (prizeNum - q[qIdx].peek() > K) { // q의 맨 앞 친구와 K 이상 차이나는 경우
                    q[qIdx].poll();
                } else { // q 안에 좋은 친구만 남은 경우
                    break;
                }
            }
            result += q[qIdx].size();
            q[qIdx].offer(prizeNum);
        }
    }
}