[백준] 좋은 친구 (자바 풀이)

 

 

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

 

 

 

📌 문제

상근이는 환갑을 바라보던 나이에 수능 시험을 다시보고 교대에 입학했고, 초등학교 선생님으로 취직했다.

• 상근: 요즘 애들은 친구를 사귀지 않나봐. 내가 앞에서 보고 있으면, 친구가 있는 학생이 별로 없는 것 같아.
• ??: 오빠! 오빠는 말콤의 친구와 성적이라는 책 안 읽어 봤어? 이 책에는 성적과 친구가 무슨 관계가 있는지 나와. 요즘 애들은 친구를 사귀기 전에 먼저 그 친구의 반 등수를 살펴봐. 말콤은 이 연구를 하기 위해서 6년동안 초등학교에서 선생님으로 위장 했었지. 하지만, 6년이라는 시간을 초등학교에서 보냈지만, 그 사람은 결국 결론을 얻지 못했어.
• 상근: 근데?
• ??: 말콤이 어느 날 자신이 초등학생이 되어 학교를 활보하는 꿈을 꾸었어. 근데 잠을 깨고 나니 내가 꿈을 꾸고 초등학생이 된건지, 아니면 초등학생이 꿈을 꾸고 지금의 내가 되어있는지를 모르겠는거야. 그래서 말콤은 상식적인 사고 방식에 큰 의문을 가졌지. 그 때 말콤은 깨달았던거야. 초등학교 친구는 부질없구나. 그제서야 알게된거야. 모든 학생은 자신과 반 등수의 차이가 K를 넘으면 친구가 아니라는거.
• 상근: 아? 근데 K는 어떻게 구해?
• ??: K는 문제에서 주어지지. 근데, 더 중요한 사실이 있어. 친구와 좋은 친구의 차이야. 말콤이 친구와 성적을 쓰고 2년 뒤에 낸 책인 좋은 친구라는 책에는 좋은 친구는 이름의 길이가 같아야 된다는 말이 나와.
• 상근: 아! 그럼 난 오늘 집에 가서 우리 반에 좋은 친구가 몇 쌍이나 있는지 구해봐야 겠어!

상근이네 반의 N명 학생들의 이름이 성적순으로 주어졌을 때, 좋은 친구가 몇 쌍이나 있는지 구하는 프로그램을 작성하시오. 좋은 친구는 등수의 차이가 K보다 작거나 같으면서 이름의 길이가 같은 친구이다.

⚔ 입력 

 

입력 

 

첫째 줄에 N과 K가 주어진다. (3 ≤ N ≤ 300,000, 1 ≤ K ≤ N) 다음 N개 줄에는 상근이네 반 학생의 이름이 성적순으로 주어진다. 이름은 알파벳 대문자로 이루어져 있고, 2글자 ~ 20글자이다.

 

출력 

 

첫째 줄에 좋은 친구가 몇 쌍이 있는지 출력한다.

 

예제 입출력

4 2
IVA
IVO
ANA
TOM

 

5

 

---

 

 

👀 문제 해석

이 문제는 학생들의 이름과 성적을 고려하여 "좋은 친구" 쌍을 찾는 문제이다. 학생의 이름이 성적순으로 주어지고, 두 학생이 "좋은 친구"가 되려면 다음 두 가지 조건을 만족해야 한다.

1. 이름의 길이가 같아야 한다.
2. 성적 등수의 차이가 주어진 정수 ( K ) 이하여야 한다.

문제에서 학생 수 ( N )은 최대 300,000명까지 주어질 수 있고, 이름의 길이는 2글자에서 20글자 사이이다. 문제의 크기와 제약 조건이 꽤 크다. 따라서 시간 복잡도가 중요한 요소인 문제라고 볼 수 있겠다. 가장 쉽게 생각할 수 있는 완전 탐색, 즉 단순히 모든 학생 쌍을 비교한다면 O(( N^2 )) 복잡도로, 시간 초과가 발생할 것이 예상된다.

 

따라서, 이 문제는 효율적인 자료구조와 알고리즘을 필요로 한다. 슬라이딩 윈도우(Sliding Window) 또는 큐를 사용한 범위 검사 문제로 분류할 수 있다. 슬라이딩 윈도우 기법을 사용하여 등수의 차이를 제한하는 범위 내에서 이름 길이가 같은 학생들을 효율적으로 관리하는 것이 문제 풀이의 핵심이다.

 

메커니즘 자체는 비슷한데, 본 풀이에서는 큐를 사용해서 풀어보고자 한다.

🔎 접근

다음과 같은 사고의 흐름을 따라 접근할 수 있다.

 

먼저, "좋은 친구"가 되기 위한 조건 두 가지를 다시 생각해보자. 이름의 길이가 같은 학생들만을 비교해야 하므로, 이름 길이에 따라 학생들을 그룹화할 필요가 있다. 또한, 등수 차이가 ( K ) 이하여야 하므로, 특정 이름 길이를 가진 학생들 사이에서 등수 차이를 검사해야 한다.

여기서 이 그룹화 아이디어가 사실 문제 풀이의 절반이기도 할 만큼 중요한 아이디어라고 생각한다.

 

문제를 풀기 위한 구체적인 아이디어를 다음과 같은 질문들을 해보면서 발전시켜보자.

• 어떻게 하면 이름 길이가 같은 학생들을 효율적으로 관리할 수 있을까?
  • 이름의 길이는 2글자에서 20글자까지로 제한되어 있다. 따라서 이를 이용해보자. 이름 길이에 따라 19개의 그룹을 만들 수 있을 것이다. 각각의 그룹은 해당 이름 길이를 가진 학생들의 리스트가 된다.
• 성적 순서대로 주어진 학생들 중, 어떻게 하면 등수 차이가 ( K ) 이하인 학생 쌍을 찾을 수 있을까?
  • 여기서 슬라이딩 윈도우나 큐와 같은 자료구조를 활용하여 유효 범위 내의 학생들을 관리하자. 왜냐하면, 등수 차이라는 개념의 특성상, 특정 등수를 벗어나는 요소는 무의미한 데이터로 간주할 수 있기 때문이다. 즉, 성적 순서대로 학생을 하나씩 살펴보면서, 현재 학생보다 ( K )를 초과하는 등수 차이를 가진 학생을 큐에서 제거한다면? 그리고 남아 있는 큐의 크기를 활용한다면? 이때, 남아 있는 큐의 크기 자체가 좋은 친구 쌍이 된다는 것이 핵심이다.

결론적으로 두 아이디어, 이름 길이에 따라 학생들을 그룹화하여 관리하는 것과 큐를 이용해 등수 차이를 검사하며 유효 범위 내의 학생들만 유지한다면 해당 큐 사이즈 자체가 좋은 친구 쌍이 된다는 것이 이 문제에 있어서 핵심 아이디어이다.

 

이를 생각해내는 것이 어렵고, 일단 생각해낸다면 풀이 자체는 쉬울 것 같다. 이 두 가지 접근을 통해 문제의 제약 조건을 효과적으로 처리할 수 있고, 시간 복잡도를 O(N)으로 줄일 수 있다는 점이 중요한 포인트이다.

 

구체적인 해결 방법을 스텝별로 생각해보자.

1. 이름 길이별로 큐를 생성한다. 배열의 각 인덱스는 이름의 길이를 나타내고, 해당 인덱스에는 이름의 길이가 같은 학생들의 인덱스를 저장하는 큐가 들어간다.
2. 성적 순서대로 학생을 처리하면서, 현재 이름 길이 그룹의 큐를 검사하여 등수 차이가 ( K )를 초과하는 경우 큐에서 학생을 제거한다. 이로서 현재 학생보다 오래된 학생들 중에서 ( K )를 초과하는 등수 차이를 가지는 학생들을 제외시킬 수 있다.
3. 큐에 남아 있는 학생들의 수가 현재 학생과 "좋은 친구"가 될 수 있는 학생 수가 된다. 이 수를 결과 값에 더하고, 현재 학생의 인덱스를 큐에 추가한다.
   
   

💡 풀이 코드

1. 이름 길이별로 큐를 생성

각 이름 길이에 대해 ArrayDeque 배열을 사용해 큐를 생성한다. 배열의 인덱스는 2부터 20까지로, 각 인덱스에는 해당 길이의 이름을 가진 학생들의 인덱스를 저장한다.

   ArrayDeque<Integer>[] queues = new ArrayDeque[21];
   for (int i = 2; i <= 20; i++) {
       queues[i] = new ArrayDeque<>();
   }

 

2. 성적 순서대로 학생을 처리하며 좋은 친구 쌍 계산

위의 로직에 따라 queues는 다음과 같이 초기화되어 있을 것이다.

 

queues[2]:  []
queues[3]:  []
queues[4]:  []
queues[5]:  []
...
queues[20]: []

각 인덱스에는 ArrayDeque<Integer>가 들어 있으며, 현재 모든 큐는 빈 상태이다. 예를 들어 queues[3]은 이름의 길이가 3인 학생들의 인덱스를 저장하는 큐가 된다.

 

이제 학생들을 순회하면서, 먼저 현재 학생의 이름 길이에 해당하는 큐를 선택한다. 예를 들어, 첫 번째 학생의 이름이 "IVA"이고, 이 이름의 길이는 3이다. 따라서 queues[3]을 선택한다. 이 큐에는 현재까지 이름 길이가 3인 학생들의 인덱스가 저장되고 있을 것이다.

 

만약에 현재 학생의 인덱스를 i라고 할 때, queues[nameLength]의 맨 앞에 있는 값과 현재 인덱스 i의 차이를 계산하여 ( K )보다 크면 등수를 벗어났다는 의미일 것이다. 따라서 이 요소에 대해서 큐에서 제거한다. 즉, 등수 차이가 ( K )를 초과하는 학생을 제거하는 것이다.

 

   while (!queues[nameLength].isEmpty() && i - queues[nameLength].peekFirst() > 등수차이) {
       queues[nameLength].pollFirst();
   }

 

여기서 peekFirst()는 큐의 맨 앞에 있는 학생의 인덱스를 가져온다. 즉, 맨 앞이므로 가장 먼 등수이다. pollFirst()로 큐에서 제거한다. 이렇게 큐의 맨 앞에 있는 학생과 현재 학생의 등수 차이가 ( K )를 초과하는 학생은 "좋은 친구"의 조건을 만족하지 않으므로 큐에서 제거한다.

 

이제, 현재 큐의 크기를 좋은 친구 쌍의 수에 더한다. 현재 유효 범위 내에 있는 학생들은 큐에 남아 있으므로, 큐의 크기가 곧 좋은 친구 쌍이 된다. 왜 그럴까?

 

큐에는 이미 이름 길이가 같은 학생들만이 추가되고, 앞서 등수 차이가 ( K )를 초과하는 학생들은 제거했기 때문에, 현재 큐에 남아 있는 학생들은 모두 "좋은 친구"의 조건을 만족할 것이다. 예를 들어보자.

 

( K = 2 )이고, 이름의 길이가 3인 학생들의 인덱스가 [0, 1, 2]라고 가정해보자.

1. 첫 번째 학생(인덱스 0)의 이름은 "IVA"이고, 큐는 처음에 비어 있다. 따라서, "IVA"를 큐에 추가하면 큐의 상태는 다음과 같다.
2. queues[3]: [0]
3. 두 번째 학생(인덱스 1)의 이름은 "IVO"이고, 현재 큐에는 [0]이 들어 있다. 인덱스 차이는 ( 1 - 0 = 1 ), 이는 ( K = 2 ) 이하이므로 큐의 맨 앞에 있는 학생을 제거하지 않고, "IVO"를 큐에 추가한다. 큐의 상태는이때, 큐의 크기가 2이므로 좋은 친구 쌍의 수에 1을 더한다. 이는 현재 학생이 인덱스 0과 좋은 친구가 될 수 있음을 의미한다.
4. queues[3]: [0, 1]
5. 세 번째 학생(인덱스 2)의 이름은 "ANA"이다. 현재 큐에는 [0, 1]이 들어 있고, 인덱스 차이는 ( 2 - 0 = 2 ), 이는 ( K )와 같으므로 큐의 맨 앞에 있는 학생을 제거하지 않는다. "ANA"를 큐에 추가하면 큐의 상태는현재 큐의 크기는 3이므로, 좋은 친구 쌍의 수에 2를 더한다. 이는 인덱스 2의 학생이 인덱스 0과 1의 학생과 각각 좋은 친구가 될 수 있다는 뜻이다.
6. queues[3]: [0, 1, 2]

이렇게 큐에 남아 있는 학생들은 모두 현재 학생과의 등수 차이가 ( K ) 이하인 학생들만 유지된다. 따라서 큐의 크기가 곧 좋은 친구 쌍의 수가 된다.

goodFriendCount += queues[nameLength].size();
queues[nameLength].offerLast(i);

마지막으로 현재 학생의 인덱스를 큐에 추가하는 코드를 작성한다. 이렇게 함으로써 이 다음에 오는 요소들에 대해서도 동일하게 등수를 판단할 수 있을 것이다.

   queues[nameLength].offerLast(i);

전체 코드는 다음과 같다.

    private static long getGoodFriends(String[] names, int 학생수, int 등수차이) {
        ArrayDeque<Integer>[] queues = new ArrayDeque[21];
        for (int i = 2; i <= 20; i++) {
            queues[i] = new ArrayDeque<>();
        }

        long goodFriendCount = 0;

        for (int i = 0; i < 학생수; i++) {
            int nameLength = names[i].length();

            // 큐의 맨 앞에 있는 학생과 등수 차이가 K보다 크면 큐에서 제거
            while (!queues[nameLength].isEmpty() && i - queues[nameLength].peekFirst() > 등수차이) {
                queues[nameLength].pollFirst();
            }

            // 현재 큐의 크기만큼 좋은 친구 쌍이 가능
            goodFriendCount += queues[nameLength].size();

            // 현재 학생의 등수를 큐에 추가
            queues[nameLength].offerLast(i);
        }

        // 좋은 친구 쌍의 수 출력
        return goodFriendCount;
    }

이제, 예제 입력 4 2와 학생 이름 목록 ["IVA", "IVO", "ANA", "TOM"]을 처리하면서 작동 과정을 따라가보자.

초기 상태

• 모든 큐는 비어 있다.

  queues[2]:  []
  queues[3]:  []
  queues[4]:  []
  ...
  queues[20]: []

첫 번째 학생 "IVA" (인덱스 0)

• 이름의 길이는 3이므로 queues[3]을 사용한다.
• queues[3]은 비어 있으므로 while 반복문을 수행하지 않고, 큐에서 제거되는 요소는 없다.
• 현재 큐의 크기는 0이므로 좋은 친구 쌍의 수에 아무것도 추가하지 않는다.
• queues[3]에 0을 추가한다.
• queues[3]: [0] goodFriendCount = 0

두 번째 학생 "IVO" (인덱스 1)

• 이름의 길이는 3이므로 queues[3]을 사용한다.
• queues[3]에는 [0]이 들어 있다. 인덱스 차이는 ( 1 - 0 = 1 ), 이는 ( K = 2 ) 이하이므로 제거하지 않는다.
• 현재 큐의 크기는 1이므로 goodFriendCount에 1을 더해 총 1이 된다.
• queues[3]에 1을 추가한다.
• queues[3]: [0, 1] goodFriendCount = 1

세 번째 학생 "ANA" (인덱스 2)

• 이름의 길이는 3이므로 queues[3]을 사용한다.
• queues[3]에는 [0, 1]이 들어 있다. 인덱스 차이는 ( 2 - 0 = 2 ), 이는 ( K )와 같으므로 제거하지 않는다.
• 현재 큐의 크기는 2이므로 goodFriendCount에 2를 더해 총 3이 된다.
• queues[3]에 2를 추가한다.
• queues[3]: [0, 1, 2] goodFriendCount = 3

네 번째 학생 "TOM" (인덱스 3)

• 이름의 길이는 3이므로 queues[3]을 사용한다.
• queues[3]에는 [0, 1, 2]가 들어 있다. 인덱스 차이는 ( 3 - 0 = 3 ), 이는 ( K )를 초과하므로 0을 큐에서 제거한다.
• 이제 queues[3]에는 [1, 2]가 남아 있고, 인덱스 차이는 각각 2와 1로 ( K ) 이하이다.
• 현재 큐의 크기는 2이므로 goodFriendCount에 2를 더해 총 5가 된다.
• queues[3]에 3을 추가한다.
• queues[3]: [1, 2, 3] goodFriendCount = 5

최종 결과

• 모든 학생을 처리한 후 goodFriendCount는 5가 된다.
• 따라서, 출력 결과는 5이다.
• 시간 복잡도는 O(N)으로, N이 최대 300,000일 때도 효율적으로 해결 가능.