[Python] 백준 3078번 좋은 친구

3078번 (좋은 친구)

queue 를 이용해서 풀었다. 예를 들어 K=2 면, 3명 단위로 블럭을 만들어 한칸씩 이동시킨다고 생각하면 된다.

처음 푼 방식

from sys import stdin
from collections import deque
N, K = map(int, stdin.readline().split())
grade = [0] * 21
student = deque()
count = 0
for _ in range(N):
    a = len(stdin.readline().rstrip())
    student.append(a)
    grade[a] += 1
    if len(student) > 0:
        if len(student) > K+1:
            old = student.popleft()
            grade[old]-=1
        count += (grade[a]-1)
print(count)

좀 더 효율적인 방식

from sys import stdin
from collections import deque
n,k = map(int, stdin.readline().split())
student = [0]*n
num = [0]*21
count = 0

for i in range(n):
    student[i] = len(stdin.readline().rstrip())
    if i > k:
        num[student[i-k-1]]-=1
    count += num[student[i]]
    num[student[i]] += 1

print(count)