시간초과를 대비하는 빠른 입력 받기 방법 sys.stdin.readline(), strip()

input()대신 sys.stdin.readline()을 사용하는 이유는?

단순히 몇줄 입력받는 문제들과는다르게, 반복문으로 여러줄을 연속적으로 입력받아야 하는(정렬, 이진 탐색, 최단 경로 문제)의 경우 input()으로 데이터를 입력받으면 시간초과로 오답판정을 받는 일이 발생할 수 있다.
따라서 입력을 최대한 빠르게 받는 방법을 알아야 한다.

sys라이브러리에 정의되어 있는 sys.stdin.readline() 함수를 이용한다.

그 대표적인 예시가 백준 BOJ 15552번 문제이다.

import sys  #sys 라이브러리를 import 

n = int(input())	# case개수를 입력 받기
for i in range(n):	# case개수만큼 반복문 실행
    a, b = map(int, sys.stdin.readline().split())	# sys.stdin.readline()으로 빠르게 데이터 입력받기
    print(a+b)

 

sys.stdin.readline() 사용법

• 한 개의 정수를 입력받을 때

import sys
a = int(sys.stdin.readline())

 

※ 그냥 a = sys.stdin.readline() 하면 어떻게 될까?

sys.stdin.readline()은 데이터를 한줄 단위로 입력받기 때문에, 개행문자가 같이 입력받아 진다.
만약, 2를 입력했다면, 3\n 이 저장되기 때문에, 개행문자를 제거해야 한다.
개행문자를 제거하려면, .rstrip()을 붙여주면 된다.

import sys

sys.stdin.readline().rstrip

하지만, 이경우에도 우리가 받고싶은 데이터 형태인 정수형이 아니게 된다.
따라서 int형을 변환을 해주어야 하는데, 처음부터 데이터를 입력받음과 동시에 int형으로 형변환을 해주면, .rstrip()을 붙여주지 않아도, 개행문자가 제거된다.

따라서 정수형으로 입력받고 싶을 때는, 아래 코드를 기억하자.

import sys

a = int(sys.stdin.readline())

 

• 정해진 개수의 정수를 한줄에 입력받을 때

import sys
a, b, c = map(int, sys.stdin.readline().split())

map() 함수는 반복 가능한 객체(리스트 등)에 대해 각각의 요소들을 지정된 함수로 처리해주는 함수이다.
위와 같이 사용하면 a, b, c에 대해 각각 int 형으로 형변환을 할 수 있다.

 

• 임의의 개수의 정수를 한줄에 입력받아 리스트에 저장할 때

import sys
data = list(map(int, sys.stdin.readline().split()))

split()은 문자열을 나눠주는 함수이다. 
괄호 안에 특정 값을 넣어주면, 그 값을 기준으로 문자열을 나누고, 아무 값도 넣어주지 않으면 공백(스페이스, 탭 엔터 등)를 기준으로 나눈다.

list()는 자료형을 리스트형으로 변환해주는 함수이다.

 

• 임의의 개수의 정수를 n줄 입력 받아 2차원 리스트에 저장하는 경우

import sys
data = []		# 빈 리스트 생성
n = int(sys.sydin.readline())	# 입력받을 데이터는 몇줄?

for i in range(n):		# 데이터 입력 받기를 n줄번 만큼 반복
	data.append(list(map(int, sys.stdin.readline().split())))

 

• 문자열 n줄을 입력받아 리스트에 저장할 때

import sys
n = int(sys.stdin.readline())
data = [sys.stdin.readline().strip() for i in range(n)]

입력

3
HELLO
I AM
MINGYO

출력

['HELLO', 'I AM', 'MINGYO']

 

 

※ lstrip(), strip(), rstrip() 사용법 및 차이점

 

공백(white space) 제거

다음 코드는 strip()에 인자를 전달하지 않는다. 인자를 전달하지 않으면 문자열에서 공백을 제거한다

[입력]

text = ' Water boils at 100 degrees '
print('[' + text.rstrip() + ']')
print('[' + text.lstrip() + ']')
print('[' + text.strip() + ']')

[출력]

[ Water boils at 100 degrees]
[Water boils at 100 degrees ]
[Water boils at 100 degrees]

 

 

동일한 문자 제거

인자로 문자 1개를 전달하면 그 문자와 동일한 것을 모두 제거한다.
동일하지 않은 문자가 나올 떄까지 제거한다.

[입력]

text = '0000000Water boils at 100 degrees 000'
print(text.lstrip('0'))
print(text.rstrip('0'))
print(text.strip('0'))

[출력]

Water boils at 100 degrees 000
0000000Water boils at 100 degrees
Water boils at 100 degrees

 

여러 문자 제거

인자로 여러 문자를 전달하면 그 문자들과 동일한 것들을 모두 제거한다. 동일하지 않은 문자가 나올 때까지 제거한다.

[입력]

text = ",,,,,123.....water....pp"
print(text.lstrip(',123.p'))
print(text.rstrip(',123.p'))
print(text.strip(',123.p'))

[결과]

water....pp
,,,,,123.....water
water

다음 예제도 위처럼 동일하게 여러 문자를 인자로 전달했다.

[입력]

text = ' Water boils at 100 degrees '
print(text.lstrip(' Water'))
print(text.rstrip(' degrees '))
print(text.strip(' degrees '))

[결과]

boils at 100 degrees
 Water boils at 100
Water boils at 100

sys.stdin.readline()