[Python] 백준 / gold / 3078 : 좋은 친구

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

슬라이딩 윈도우 문제. (3 ≤ N ≤ 300,000, 1 ≤ K ≤ N) 이기 때문에 나이브한 알고리즘으로는 O(N * K) 가 10^10 을 넘어갈 수 있다. 그래서 부르트 포스로는 풀 수 없고, K 가 불변값이기 때문에 슬라이딩 윈도우로 푸는게 적절했던거 같다. 슬라이딩 윈도우로는 O(N + K) 소요되기 때문에 시간제한을 통과할 수 있다.

collections 의 default dict 를 사용했는데 defaultdict(int) 이 문법을 기억해야겠다.. 안에 0을 넣으면 에러가 난다. 값이 아닌 '타입'을 넣어줘야 한다!

좋은 친구의 "쌍"을 구하는 거라서 1씩 더해주는 것이 아니고 combination 연산이다. answer += window[nameLen[i]] . 현재 윈도우에 있는 친구들에게 새로 들어올 친구를 한번씩 짝지어 줄 수 있으므로 이렇게 연산하면 combination 을 세주는 효과를 낼 수 있었다.

파이썬 코드

import sys
from collections import defaultdict
N, K = map(int, sys.stdin.readline().split())
answer = 0
nameLen = []

for _ in range(N):
    nameLen.append(len(sys.stdin.readline().strip()))

window = defaultdict(int)
# window 초기 세팅
for i in range(K+1):
    if window[nameLen[i]] > 0:
        answer += window[nameLen[i]]
    window[nameLen[i]] += 1

for i in range(K+1, N):
    window[nameLen[i-K-1]] -= 1

    if window[nameLen[i]] > 0:
        answer += window[nameLen[i]]
    window[nameLen[i]] += 1

print(answer)