Algorithm/BOJ

파이썬으로 백준 문제 풀 때, 여러 개 입력 받기 tip

강승원 2024. 1. 10. 19:47

map을 유용하게 쓰는 것 같다.

1.  한 줄로 받기

example.

1 2 3 4 5 6

이런 식으로 한 줄로 들어오는 숫자를 리스트로 받고 싶다. 근데 숫자니까 이것들을 int형으로 형변환을 하고싶다. 그러면 map을 쓰면 된다. map은 파이썬에서 제공하지만 사실 js에도 있고 다른 언어에도 있는 프로그래밍 언어 국룰! 순회 가능한 자료형의 메소드 혹은 함수다. js는 메소드로 제공되지만 파이썬은 함수형태이고 

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

이렇게 써준다. 

첫번째 파라미터는 함수(여기서 int는 형변환 해줄 때 사용하는 int)
두번째 파라미터는 순회 가능한 자료형(튜플, 리스트) 를 넣어주면, 두번째 파라미터에 넣은 자료형 원소마다 첫번째 파라미터에 넣은 함수를 실행시켜준다.

그리고 나온 이 input_list 를 보면 당황스러울 수 있는게 한가지 있는데 이 친구를 list로 착각하고 slicing을 하다가 에러를 만나게 된다. 

왜냐하면, 이때 map에서 나온 친구는 list가 아닌 iterator 객체이기 때문이다. 

map은 iterable한 변수를 받아 iterator객체를 반환해주는 함수!

iterator에 대해선 나중에 자세히 다루어 보도록 하겠다.

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

아무튼 그래서 list로 한번 더 감싸 주어야 한다.

리스트가 아닌 형태로 받고 싶을 때

1열로 들어오는 숫자를 그냥 변수에 할당하고 싶을 수 있다.

1 2 3

이 들어올때, 1 은 N 2는 M 3은 K 그럴 땐, 

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

쉼표로 나누어서 받으면 된다.

2. 여러줄로 들어오는 경우

1
2
3
4
5

이렇게 들어온다고 했을 때

N = int(input())
input_list =[]
for _ in range(N):
	input_list.append(int(input()))

초보자 시절 다 이렇게 했을 것이다.

하지만 한줄로 섹시하게 처리하는 방법이 있다

input_list = [int(input()) for _ in range(N)]

바로

List Comprehension 을 사용하는 것. set도 이렇게 만들 수 있다.

input_list = {int(input()) for _ in range(N)}

괄호를 바꿔주면 중복이 제거된 set이 된다.

 

3. 두 줄 받고 합치고 싶을때

1 2 3 4 5
a b c d e

위 두 개의 입력을 리스트로 받아서 순차적으로 묶고 싶다면?

1. tuple

number_list = list(map(int,input().split()))
string_list = input().split()
input_tuple = [(i,j) for i in number_list for j in string_list]
//zip을 써도 된다
input_tuple = list(zip(number_list,string_list))
// list로 감싸는 이유는 zip만 쓰면 zip객체가 되어버리기 때문이다

 

이렇게 한다면 튜플의 쌍으로 묶인 리스트가 완성!

2. dictionary

key value 쌍의 dictionary가 되고 싶다면?

input_dict = {key: value for key, value in zip(string_list, number_list)}

 

 

4. 바둑판을 만들고 싶을때(2차원 배열 바로 넣기)

가끔 바둑판이나 체스판으로 머리 아프게 하는 문제가 있다. 예를 들어 아래 문제와 같은 경우

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

이런 형태의 인풋은

grid = [list(map(str,input())) for _ in range(N)]

이렇게 처리 해준다. 위에 설명했던 comprehension이랑 map을 합쳐서 썼다. 이러면 복잡하게 2차원 배열을 만들 필요가 없다.

5. 번외

input()  시간초과 날 때, 이 경우는 아마 검색하면 많이 나올텐데 자세한건 나중에 조사해보고

sys 모듈의 stdin readline을 사용하면 해결이 된다더라

import sys
input = sys.stdin.readline
input()

대부분 사람들이 이렇게 쓰고 있는데, 이러면 문제가 가끔 특정 경우에  끝 쪽에 빈 공백이 들어간다.

그래서 필자는

import sys
input = lambda:sys.stdin.readline().strip()

strip을 해준 lambda함수를 변수에 할당해준다.  그럼 평소의 input처럼 사용할 수 있다.

요즘은 이게 습관이 되서 그냥 문제 보고 위에 코드를 바로 입력해주고 있다.

추가로 더 알게 된다면 이 글에 추가하도록 하겠다.