[Python][백준] 3078번 좋은 친구

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

공부 날짜 : 2023.01.31
정답 참조 여부 : X

학생이 순서대로 주어질때 k범위 이내의 학생중 이름의 길이가 같은 학생의 수를 구하는 문제이다.

---

가장 정석인 방법은 i번째 학생을 기준으로 k만큼 계속 학생이름의 수를 카운트 하면 되겠지만 그러면 O(N*M)이기 때문에 시간초과가 발생한다.

그래서 떠올린 방법은 k범위의 학생들의 이름길이를 리스트로 하고 학생의 수를 저장해두고, i를 옮길때 만다 i+k번째 학생의 이름을 갱신하고 i번째 학생을 제거하여 길이가 같은 학생의 수를 result에 더해주는 방식으로 구했다.

스택의 방법만 생각하면서 문제를 풀고 있었는데 이게 자료구조가 맞나? 생각이 들어서 그냥 내 방식으로 풀어서 정답으로 나왔다.

알고리즘을 확인하니 스택이 아니라 큐 구조였고 큐를 이런식으로 구현할 수 있음을 알 수 있는 문제였다.

소스코드

import sys
input_ = sys.stdin.readline
##########################################
n, m = map(int, input().split())

students = [len(input_().rstrip()) for _ in range(n)]

# k범위 이내에서 학생이름의 길이가 같은 학생의 수
# 이름길이가 3인 학생들의수가 len_list[3]에 저장됨
len_list = [0] * 21

# 초기 1~m번의 범위 학생들 수 이름별로 정리
for i in range(m):
    len_list[students[i]] += 1

result = 0

# n-m번까지에서 좋은 친구의 수 찾기
# n-m번 부터는 범위가 m보다 작아짐
for i in range(n-m):
    j = i+m
    len_list[students[j]] += 1
    len_list[students[i]] -= 1
    result += len_list[students[i]]

# n-m번 학생부터 나머지 학생들간의 좋은 친구 수 찾기
for i in range(n-m, n):
    len_list[students[i]] -= 1
    result += len_list[students[i]]

print(result)