[Sliding Window] 백준 3078번 “좋은 친구” Python 풀이

문제 링크

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

풀이

• 문제에 대한 설명이 길지만, 결국 등수의 차이가 K보다 작거나 같으면서 이름의 길이가 같은 친구 쌍이 얼마나 있는지를 찾는 문제입니다.
• 해당 문제는 슬라이딩 윈도우 알고리즘을 사용해 해결할 수 있습니다.
• 딕셔너리를 사용해, k 크기를 가지는 윈도우 내에 존재하는 학생들을 {이름 길이 : 인원} 형식으로 저장합니다.
• 윈도우를 오른쪽으로 한 칸씩 이동하며,
• 윈도우에서 빠지는 학생은 딕셔너리에서 제거하고,
• 윈도우에 추가되는 학생은 딕셔너리에 추가합니다.
  • 이때, 추가되는 학생의 이름 길이와 동일한 길이가 딕셔너리에 존재한다면, 친구 쌍의 개수를 +1 해줍니다.

코드

import sys
from collections import defaultdict

n, k = map(int, sys.stdin.readline().split())

# 성적순으로 나열된 학생의 이름 길이를 저장
arr = []
for _ in range(n):
    arr.append(len(sys.stdin.readline().rstrip()))

# 윈도우 내 학생들에 대한 {이름 길이 : 인원} 저장
dict = defaultdict(int)

# 친구 쌍의 수
count = 0

for end in range(n):

    # 윈도우를 오른쪽으로 한 칸 씩 이동
    if (end > k):
        # 구간에서 빠지는 학생에 대한 이름 길이 인원을 감소
        dict[arr[end-k-1]] -= 1

    # 구간에 추가되는 학생의 이름 길이와 동일한 이름 길이를 가진 학생이 있다면, 친구 쌍의 수 증가
    count += dict[arr[end]]

    # 구간에 추가되는 학생에 대한 이름 길이 인원을 추가
    dict[arr[end]] += 1

print(count)