[BOJ/PYTHON] 3078. 좋은 친구

백준 문제 링크: https://www.acmicpc.net/problem/3078

문제 요약

좋은 친구가 몇 쌍 있는지 구하기

모든 학생은 자신과 반 등수의 차이가 K를 넘으면 친구가 아님

좋은 친구는 이름의 길이가 같아야 함

-> 좋은 친구는 등수의 차이가 K를 넘지 않으며 이름의 길이가 같음

핵심 아이디어

이름 길이에 해당하는 딕셔너리를 사용한다. 키값은 길이, value값은 그 길이인 이름 개수이다.

이름을 순서대로 보면서 일단 큐에 이름을 넣고 딕셔너리에 그 이름 길이에 해당하는 키값에 1을 더해주면 된다.

확인하는 사람의 범위가 k를 넘어가면 앞에서 하나씩 빼준다. 그러면 등수의 차이가 k 이내를 계속 유지하게 된다. 슬라이딩 윈도우를 덱으로 구현한 셈이다.

해당 딕셔너리에서 이름 길이가 같은 게 있으면 좋은 친구 카운트를 더해준다. 이름 길이가 같은 경우가 없다면 딕셔너리 값이 0이고 0을 더해주는 건 의미가 없기 때문에 따로 분기처리를 안 해줘도 된다.

좋은 친구 카운트를 더해줄 땐 이번에 확인할 사람을 큐에 넣어주기 전에 먼저 처리해줘야 한다.

풀이

# gold4-3078. 좋은 친구

import sys
from collections import deque
input = sys.stdin.readline

n, k = map(int, input().split())
len_names = {i : 0 for i in range(2, 21)} # 2 ~ 20글자
q = deque()
answer = 0

for i in range(n):
    name = input().rstrip() # 성적 순으로 입력됨
    if i > k: # 범위 넘어가면 하나씩 빼줌
        popped_name = q.popleft()
        len_names[len(popped_name)] -= 1 # 딕셔너리에서도 빼줌
    answer += len_names[len(name)] # 이름 같은 수만큼 더함
    q.append(name)
    len_names[len(name)] += 1 # 딕셔너리에 더해줌

print(answer)

실행 시간

메모리 42904KB, 시간 304ms (python3)