[Python] 한 번에 여러 값 입력 (map 함수 / split 함수)

알고리즘 문제를 풀다가 여러 값을 입력받았을 때 어떤 식으로 해결하면 좋을지 고민하게 됐다. 이전에 백준에서 문제를 풀 때 대충 찾아봐서 해결을 하긴 했지만 왜 그렇게 써야 하는지 제대로 이해하지 못한 거 같아 이번에 정리를 해보려고 한다.

1. 값 입력

python에서 어떤 값을 입력받고 싶을 때 input() 함수를 사용한다. input만 단독으로 쓰면 문자열로 저장을 하고, int() 또는 float()를 사용하면 각각 정수와 실수로 저장한다.

a = input() # 문자열로 저장
b = int(input()) #정수로 저장
c = float(input()) # 실수로 저장

 

2. 여러 값 입력 (split)

또한, 한 번에 여러 값을 입력하고 싶을 땐 split()을 사용할 수 있고, 각 값을 공백을 기준으로 구분하여 순서대로 저장한다. 입력된 값은 리스트의 형태로 저장된다.

a = input().split()

만약 내가 hello world를 중간에 공백을 두고 입력하게 된다면 아래와 같은 결과를 보여준다.

>>> ['hello', 'world']

hello world 대신에 숫자 3과 5를 입력하면 결과는 다음과 같다

>>> ['3', '5']

위에서 말했다시피 input함수만 쓰면 문자열로 저장하기 때문에 숫자를 입력하더라도 문자열로 그 값을 반환한다. 받은 값을 정수로 바꾸려면 int()를 사용하면 된다.

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

 

3. map() 함수

하지만 계속해서 int를 쓰는 것은 귀찮은 과정이다. 이럴 때  map() 함수를 사용할 수 있다. map은 여러 개의 데이터를 한 번에 다른 형태로 바꿀 때 사용된다.

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

위 코드와 같이 쓰면 여러 값을 입력할 수 있고 그 값을 정수로 바꿀 수 있다. 실수로 바꾸고 싶을 땐 int 대신에 float를 쓰면 된다. 

또한 입력된 여러 값을 리스트의 형태로 저장할 수도 있다. 

num_list = list(map(int, input().split()))
print(num_list)

위에서 본 map 함수를 list로 감싸면 입력한 값들을 리스트에 저장할 수 있다. 예를 들어 1에서 5까지 입력을 때 결과는 다음과 같다.

>>> [1, 2, 3, 4, 5]

 

4. sys.stdin.readline()

이 함수도 여러 개의 값을 받을 수 있다는 점은 input().split()과 비슷하지만, 반복문으로 여러 줄을 입력받는 경우에는 sys.stdin.readline()을 사용해야 시간 초과가 발생하지 않는다.

우선 이 함수를 쓰기 위해서는 import sys를 해 줘야 한다. 한 개의 값을 받을 때 아래와 같이 쓸 수 있다.

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

그냥 sys.stdin.readline()만 사용하면 변수가 문자열로 저장되기 때문에 정수로 바꾸기 위해서는 앞에 int를 사용해야 한다. 또한 이 함수는 데이터를 한 줄 단위로 입력받기 때문에 개행 문자가 같이 입력된다.

개행 문자를 제거하기 위해서는 .rstip()을 사용할 수도 있다.

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

이 경우에는 테이터가 정수가 아니게 된다. 데이터를 정수로 바꿔주면서 개행문자를 제거하려면 int(sys.stdin.readline()) 이라고 쓰면 된다.

 

또한 여러 값을 받을 때에는 input()의 예시처럼 map()과 split()을 사용할 수 있다.

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

 

값들을 리스트에 저장할 때에는 list()를 사용하면 된다.

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