[백준/Baekjoon] 3078번: 좋은 친구

1. 문제

상근이는 환갑을 바라보던 나이에 수능 시험을 다시보고 교대에 입학했고, 초등학교 선생님으로 취직했다.

• 상근: 요즘 애들은 친구를 사귀지 않나봐. 내가 앞에서 보고 있으면, 친구가 있는 학생이 별로 없는 것 같아.
• ??: 오빠! 오빠는 말콤의 친구와 성적이라는 책 안 읽어 봤어? 이 책에는 성적과 친구가 무슨 관계가 있는지 나와. 요즘 애들은 친구를 사귀기 전에 먼저 그 친구의 반 등수를 살펴봐. 말콤은 이 연구를 하기 위해서 6년동안 초등학교에서 선생님으로 위장 했었지. 하지만, 6년이라는 시간을 초등학교에서 보냈지만, 그 사람은 결국 결론을 얻지 못했어.
• 상근: 근데?
• ??: 내 말 아직 안끝났어. 말콤이 어느 날 자신이 초등학생이 되어 학교를 활보하는 꿈을 꾸었어. 근데 잠을 깨고 나니 내가 꿈을 꾸고 초등학생이 된건지, 아니면 초등학생이 꿈을 꾸고 지금의 내가 되어있는지를 모르겠는거야. 그래서 말콤은 상식적인 사고 방식에 큰 의문을 가졌지. 그 때 말콤은 깨달았던거야. 초등학교 친구는 부질없구나. 그제서야 알게된거야. 모든 학생은 자신과 반 등수의 차이가 K를 넘으면 친구가 아니라는거.
• 상근: 아? 근데 K는 어떻게 구해?
• ??: K는 문제에서 주어지지. 근데, 더 중요한 사실이 있어. 친구와 좋은 친구의 차이야. 말콤이 친구와 성적을 쓰고 2년 뒤에 낸 책인 좋은 친구라는 책에는 좋은 친구는 이름의 길이가 같아야 된다는 말이 나와.
• 상근: 아! 그럼 난 오늘 집에 가서 우리 반에 좋은 친구가 몇 쌍이나 있는지 구해봐야 겠어!

상근이네 반의 N명 학생들의 이름이 성적순으로 주어졌을 때, 좋은 친구가 몇 쌍이나 있는지 구하는 프로그램을 작성하시오. 좋은 친구는 등수의 차이가 K보다 작거나 같으면서 이름의 길이가 같은 친구이다.

2. 입력

첫째 줄에 N과 K가 주어진다. (3 ≤ N ≤ 300,000, 1 ≤ K ≤ N) 다음 N개 줄에는 상근이네 반 학생의 이름이 성적순으로 주어진다. 이름은 알파벳 대문자로 이루어져 있고, 2글자 ~ 20글자이다.

3. 출력

첫째 줄에 좋은 친구가 몇 쌍이 있는지 출력한다.

4. 예제 입력 1

6 3
CYNTHIA
LLOYD
STEVIE
KEVIN
MALCOLM
DABNEY

5. 예제 출력 1

2

6. 풀이

사용하는 리스트는 두 가지가 있다.

1. 학생을 성적순으로 저장하는 리스트
2. 학생의 이름의 길이에 따라 저장하는 리스트

학생을 추가하는 원리는 기본적으로 학생들을 추가하면서 이전에 이미 리스트에 존재하는 학생 중 조건에 부합하는 학생을 쌍으로 만들어 더하는 방식이다.

학생을 추가하는 중에 만약 학생의 등수가 K보다 커지면, 즉 등수의 차이가 K보다 큰 학생이 생기게 된다면 이제 자신과는 상관 없는 등수가 크게 차이나는 학생을 제거하게 된다.

for rank in range(n):
    name = len(stdin.readline().rstrip())  # 이름을 입력받음
    students[rank] = name  # 학생의 등수와 이름의 길이를 저장
    if rank > k:  # 만약 학생의 등수가 K보다 커지는 경우
        data[students[rank - k - 1]] -= 1  # 자신과 상관 없는 등수의 학생을 제거
    count += data[name]  # 자신과 이름의 길이가 같은 친구를 쌍으로 추가
    data[name] += 1  # 이름의 길이를 저장하는 리스트에 자신을 추가

또한 입력에서 시간 초과 방지를 위해 sys 모듈의 sys.stdin을 사용한다. 이유는 이 글 에서 확인할 수 있다.

7. 코드

from sys import stdin

n, k = map(int, stdin.readline().rstrip().split())
students = [0] * n
data = [0] * 21
count = 0

for rank in range(n):
    name = len(stdin.readline().rstrip())
    students[rank] = name
    if rank > k:
        data[students[rank - k - 1]] -= 1
    count += data[name]
    data[name] += 1

print(count)