백준 3078번 - 좋은 친구(C++)

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

 

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

 

단순하게 생각하면 O(N * K)로 해결할 수 있지만 최적화를 통해 O(N)으로 통과시켜야하는 문제이다.

 

AAA

BBB

CC

DDD

E

 

예제를 가정하고 K = 2라 하자.

큐의 길이가 K + 1 될 때 까지 큐에 이름의 길이를 넣고 맵에 개수를 저장한다. 

 

큐 : AAA BBB CC

맵: [3] = 2, [2] = 1

 

큐가 K + 1만큼 차면 큐에서 하나를 꺼내 꺼낸 이름의 개수가 몇개 있는지 확인하고 답을 갱신한다. 모두 친구가 될 수있는 범위만 모아놓았기 때문에 여기서는 개수만 새면 된다.

 

이렇게 하면 O(N)으로 문제 풀이가 가능하다.

 

#include <iostream>
#include <vector>
#include <map>
#include <queue>
using namespace std;

map<int, int> nums;

int main(){
    int N, M;
    cin >> N >> M;

    queue<int> q;

    for(int i = 0; i < N; i++){
        string tmp; cin >> tmp;

        q.push(tmp.size());
    }
    long long answer = 0;
    queue<int> tmpQ;
    while(!q.empty()){
        int now = q.front();

        if(tmpQ.size() < M + 1){
            tmpQ.push(now);
    
            nums[now]++;
            q.pop();
        }else{
            int cal = tmpQ.front();
            tmpQ.pop();

            answer += nums[cal] - 1;
            nums[cal]--;
        }
    }

    while(!tmpQ.empty()){
        int now = tmpQ.front();
        tmpQ.pop();

        answer += nums[now] - 1;
        nums[now]--;
    }

    cout << answer << '\n';
}