상근이는 환갑을 바라보던 나이에 수능 시험을 다시보고 교대에 입학했고, 초등학교 선생님으로 취직했다.
상근이네 반의 N명 학생들의 이름이 성적순으로 주어졌을 때, 좋은 친구가 몇 쌍이나 있는지 구하는 프로그램을 작성하시오. 좋은 친구는 등수의 차이가 K보다 작거나 같으면서 이름의 길이가 같은 친구이다.
4 2
IVA
IVO
ANA
TOM
5
6 3
CYNTHIA
LLOYD
STEVIE
KEVIN
MALCOLM
DABNEY
2
from collections import deque
if __name__ == "__main__":
N, K = map(int, input().split())
name_len = [len(input()) for _ in range(N)] # 이름의 길이만 리스트에 담는다.
cnt = 0 # 좋은 친구 수
__list = [deque() for _ in range(21)] # 이름 길이 별로 등수 저장할 리스트. 안에는 큐 구조.
# i: 등수 - 1, v: 이름 길이
for i, v in enumerate(name_len):
# 큐가 비어있지 않고, 큐[0]과 현재 학생의 등수가 K 이상 차이나면 popleft()
while __list[v] and i - __list[v][0] > K:
__list[v].popleft()
# 큐 안에 있는 등수의 학생들은 현재 학생과 좋은 친구
# 큐 안의 원소 수 만큼 cnt 증가
if __list[v]:
cnt += len(__list[v])
# 큐에 현재 학생 append
__list[v].append(i)
print(cnt)
이름의 내용을 필요가 없다. 그래서 입력받을 때 이름의 길이만 구해서 리스트에 담는다.
이름 길이 별로 등수를 담을 큐들을 리스트로 만든다.
큐가 비어있으면 학생 등수를 추가.
큐가 비어있지 않다면 큐 맨 앞 학생의 등수(index 0 원소)와 현재 학생의 등수의 차이를 구함. 만약 K보다 크면, 큐의 맨 앞 학생 제거.
위 과정을 등수차이가 K 이하이거나 큐가 빌 때까지 반복.
위 과정이 끝나면 큐 안에 있는 학생은 현재 학생과 좋은 친구. 큐 안의 길이를 cnt에 추가 후, 큐에 현재 학생 추가.