[C++] 백준 3078번 좋은 친구

문제

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

입력

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

출력

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

풀이

이름이 최대 20글자이므로 deque을 20개 만들어 이름 길이 별로 deque에 담는다. 새로운 이름이 들어왔을때 랭킹이 k이하인 친구만 남기고 나머지는 다 버린다. 이후 남아있는 친구들을 더해주고 이를 반복하면 된다

#include <iostream>
#include <vector>
#include <cstring>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <deque>
#define endl '\n'

using namespace std;

int n,k;
int ranking[300000];
deque<int> student[21];

long long solve(){
  long long ans = 0;

  for (int i=0; i<n; i++){
    deque<int> &deq = student[ranking[i]];
    deq.push_back(i);
    if (deq.size() == 1) continue;
    
    while (i - deq.front() > k){
      deq.pop_front();
    }
    ans += deq.size()-1;
  }

  return ans;
}

int main(){
  ios_base :: sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);

  // ifstream cin;
  // cin.open("input.txt");

  cin >> n >> k;

  for (int i=0; i<n; i++){
    string a;
    cin >> a;
    ranking[i] = a.size();
  }

  cout << solve();
}