[알고리즘] Python으로 다양한 입력 받기

백준에서 Python으로 알고리즘 문제풀이를 하기 위해서는 우선 문제의 입력을 받아야한다. Python으로 입력을 받는 다양한 방법을 알아보자.

input()을 이용하여 하나의 데이터 입력 받기

input()으로 입력받은 데이터는 기본적으로 string이다.

# 하나의 데이터를 "string" type으로 입력 받음
a = input()

만약 입력 받는 데이터가 정수(int)나 실수(float)라면 다음과 같이 형 변환을 할 수 있다.

# 하나의 정수(int) 데이터 입력 받기
a = int(input())

# 하나의 실수(float) 데이터 입력 받기
b = float(input())

 

두개 이상의 데이터 입력 받기

문제를 풀다 보면 다음과 같이 입력이 2개 이상인 경우가 있다.

입력이 공백으로 구분된 경우 split()을 이용하여 입력을 받을 수 있다.

# 예시 입력을 받으면 a = 6, b = 3이 저장된다.
a, b = input().split()

단, split() 위와 같은 경우 저장된 6과 3은 string 타입이다. int 타입으로 저장하기 위해서는 map 함수를 이용하여 형변환을 적용해야 한다.

a, b = map(int, input().split())

map 함수는 map(적용할 함수, iterable한 자료형)의 형태로 사용되며 iterable한 자료형의 각 원소들에 함수를 적용하고 그 결과를 map 객체로 반환한다. 위의 예시의 경우에는 두번째 인자로는 list 자료형인 ['6', '3']이 들어가고 int는 형변환을 시켜주는 함수로 첫번째 인자로 사용된다. 따라서 list의 원소인 '6'과 '3'에 각각 int 함수가 적용되어 형변환된다.

 

여러개의 데이터를 리스트로 입력 받기

그런데 공백으로 구분된 데이터의 개수가 너무 많은 경우가 있다. 아래 예제 입력의 2번째 줄 같은 경우이다.

이런 경우는 list에 한번에 저장하는 방법을 사용할 수 있다. map의 반환형은 map 객체이므로 list로 감싸서 list 타입으로 변환한다.

lst = list(map(int, input().split()))

# lst = [1, 10, 4, 9, 2, 3, 8, 5, 7, 6]

 

2차원 배열 입력 받기

위의 예제 입력은 N*M 크기의 2차원 배열로 첫번째 줄은 순서대로 N, M의 값이다. 2번째 줄부터는 배열의 1행~N행이 주어진다. 2차원 배열을 lst 변수에 저장하기 위해서는 다음과 같이 입력을 받으면 된다.

N, M = map(int, input().split())

lst = [list(map(int, input().split())) for _ in range(N)]

# lst = [[1, 1, 1], 
#        [2, 2, 2], 
#        [0, 1, 0]]

 

시간제한에 걸리는 경우

문제의 입력이 매우 많은 경우 input()을 사용하면 시간초과가 발생할 수 있다. sys.stdin.readline()을 사용하면 좀 더 빠르게 입력을 받을 수 있다. 위 예시의 input() 부분만 sys.stdin.readline()로 바꿔주고 맨위에 import sys를 추가해주면 된다.

import sys

# 정수 데이터 하나 입력 받기
a = int(sys.stdin.readline().rstrip())

# 정수 데이터 2개 입력 받기
a, b = map(int, sys.stdin.readline().split())

# 정수 데이터 리스트로 입력 받기
lst = [list(map(int, sys.stdin.readline().split())) for _ in range(N)]