백준 3078번 - 좋은 친구 - 파이썬(Python)

문제 링크

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

 

문제

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

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

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

 

입력 조건

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

출력 조건

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

 

입력 예시1

4 2
IVA
IVO
ANA
TOM

출력 예시1

5

입력 예시2

6 3
CYNTHIA
LLOYD
STEVIE
KEVIN
MALCOLM
DABNEY

출력 예시2

2

 

풀이 전 생각

 

(잘못된 생각/시간초과)

  우선, 좋은 친구의 조건은 두 가지 이다. 

1.등수의 차이가 K보다 작거나 같다

2.이름의 길이가 같다.

빈 리스트를 만들고, 입력받은 순서에 따라 리스트에 append 를 해주면 리스트의 인덱스+1 이 해당 학생의 등수가 될 것이다.

그럼 이제 반복문을 돌면서 서로 다른 두 학생이 좋은 친구인지 아닌지 확인하면 된다.

근데, 시간 제한에는 문제가 없을까?

일일히 두 학생을 비교하며 확인하게 된다면 O(NK)의 시간복잡도를 가지게 될 것이다. N의 범위가 30만까지 주어졌고, 극단적으로 N,K를 30만이라고 한다면 대략적으로는 O(NK)에서 900억 정도의 연산이 필요할 것이다.

조금 더 정확히 본다면, 실제로 비교하는 횟수는 30만명 중 두명을 고르는 경우의 수와 같기 때문에 30만 combination 2 로 450억 정도의 연산이 필요할 것이다.

그래서 다른 방법을 찾아야 할 것 같다.

 

----------------------------------------------------------------------------------------------------------------------

 입력을 전부 받고 서로 다른 두명이 좋은 친구인지 각각 확인하는 것이 아니라, 입력을 받음과 동시에 좋은 친구의 조건을 만족한다면 count를 해주는 방식으로 for 문을 N만큼 돌면서(O(N)) 그 안에서는 O(1)의 시간이 소요되는 연산을 하면 되지 않을까? 

그리고 이때 자료구조로는 list를 사용하지 않고 deque을 사용해야 할 것 같다. (아래 링크는 관련 내용에 관한 설명이다.) 

https://js-kkk.tistory.com/11

백준 2164번 - 카드2 - 파이썬(Python)

여기까지 생각을 하고 여러 풀이 방법을 찾아 보았다.

 

 

예제 코드1(PyPy3 에서는 통과, Python 3 에서는 시간 초과)

from collections import deque
n,k=map(int,input().split())
students = deque()
name_count = [0]*21
count = 0
for i in range(n): 
    if i > k: 
        name_count[students.popleft()]-=1
    name_length = len(input())
    students.append(name_length)
    name_count[name_length] += 1
    count += (name_count[name_length]-1)
print(count)

예제 코드2(둘 다 런타임 에러 (IndexError))

import sys
input = sys.stdin.readline
from collections import deque
n,k=map(int,input().split())
students = deque()
name_count = [0]*21
count = 0
for i in range(n): 
    if i > k: 
        name_count[students.popleft()]-=1
    name_length = len(input())
    students.append(name_length)
    name_count[name_length] += 1
    count += (name_count[name_length]-1)
print(count)

예제 코드3(둘 다 통과 #strip추가)

import sys
input = sys.stdin.readline
from collections import deque
n,k=map(int,input().split())
students = deque()
name_count = [0]*21
count = 0
for i in range(n): 
    if i > k: 
        name_count[students.popleft()]-=1
    name_length = len(input().strip())
    students.append(name_length)
    name_count[name_length] += 1
    count += (name_count[name_length]-1)
print(count)

 

생각

  문제를 풀면서 여러 문제점을 만났다. 처음 작성한 예제코드 1을 보면 분명 O(N)의 시간복잡도를 가지는 것 같은데 시간 초과가 나와서 당황스러웠다. 원인을 계속해서 찾던 중 아래 글을 보게 되었다.

https://ralp0217.tistory.com/entry/Python3-%EC%99%80-PyPy3-%EC%B0%A8%EC%9D%B4

Python3 와 PyPy3 차이

이 글을 보고 PyPy3 로 채점을 하니 통과를 했지만, Python 3에서는 통과가 안되는 이유가 궁금했다. N이 30만으로 극단적으로 크다고 가정해도 for문안에서의 연산들이 O(1)로 계산되면 시간은 충분할 것 같은데..  

gpt에 질문하니 input()을 사용할 때 input = sys.stdin.readline 로 사용하는 방법을 제안해줬다. input()과 sys.stdin.readline의 차이는 무엇일까? 아래는 input()과 sys.stdin.readline의 차이에 대한 간단한 설명글이다.

https://yoon001.tistory.com/74

 

또, vsc에서는 input()뒤에 strip()을 추가하든 안하든 코드가 정상적으로 동작하는 것 같은데 백준에서 채점할 때는 왜 안될까? 아래의 글과 관련이 있는 것 같다.

https://cuckoobird.tistory.com/8

 

의문을 해소하려면 vsc와 백준의 채점기계에 대해서 알아보아야 할 것 같은데.. 이것도 나중에 찾아보고 정리해야겠다