Python 숏코딩 관련 간단한 배열 입력 트릭

# 처음 주어지는 라인이 정수 N인 경우(혹은 문제 해결에 있어 필요 없는 경우)
L = list(map(int, [*open(0)][1:]))

# 첫 라인 부터 element가 주어지는 문제의 경우, unpacking이 필요없음.
L = list(map(int, open(0)))

# 이후 코드에서 index 없이 iterate 할 경우 list로 변환하지 않아도 된다.
L = map(int, [*open(0)][1:])
for e in L:
    print(e)

비교적 짧은 코드로 N개의 정수가 주어지는 문제에서 처음 주어지는 N을 제외한 eof까지 모든 라인 입력을 정수로 가져올 수 있다.

 

 

온라인 저지 사이트에서 아래 코드와 같은 동작을 한다고 생각하면 된다.

N = int(input())
L = []
for i in range(N):
    L.append(int(input()))

 

 

정렬이 필요한 경우도 다음과 같이 간략하게 쓸수 있다.

 

list로 변환할 필요없이 map 객체를 sorted()에 파라미터로 넣어도 리턴값은 list로 나온다.

L = sorted(map(int, [*open(0)][1:]))

 

 

이렇게 할 수 있는 이유는 open() 함수의 리턴값이 iteratable하기 때문이다.

 

python 공식 문서를 보면, open() 함수는 모드에 따라 리턴하는 객체의 종류가 달라지는데, 텍스트 모드에서는 io.TextIOWrapper라는 io.TextIOBase의 서브클래스를 리턴한다.

 

디폴트값은 'r' 이므로 위의 예제에서도 0, 즉 stdin을 읽기 전용 텍스트 모드로 열기 때문에 역시 io.TextIOWrapper를 리턴한다.

 

또 io.TextIOBase는 io.IOBase 추상 클래스의 서브 클래스인데, 이 추상 클래스가 Iterator Protocol을 지원하기 때문에 우리가 이 성질을 이용할수 있게 된다.

 

io — Core tools for working with streams — Python 3.10.0 documentation

 

보통 어떤 클래스에 __iter__() 와 __next__() 이 두가지의 Special Method(혹은 Magic Method)가  구현되어 있다면 이 클래스는 Iterator Protocol을 지원한다고 말할 수 있는데, 일반적으로 __iter__() 는 현재 iterator가 가리키는 값을 리턴하는 함수로, __next__() 는 iterator를 증가시키는 함수로 구현된다.

 

 

C++로 비유하자면 다음의 코드와 유사하다고 할 수 있다.

class IteratorProtocolExample
{
    ...
    
    ElementType __iter__()
    {
        return *this->iterator;
    }
    
    void __next__()
    {
    	++this->iterator;
    }
    
    Iterator<ElementType> iterator;
    
    ...
    
};

 

이 트릭으로 간단한 문제를 풀어보자.

2309번: 일곱 난쟁이

import random

L = list(map(int, open(0)))
while(1):
    s = random.sample(L, 7)
    if sum(s) == 100:
        print(*sorted(s), sep='\n')
        break

원래는 길이가 9인 배열에서 모든 순서없는 원소 쌍 (a, b)에 대해 (a, b)를 제외한 합을 구하면서 100이 되면 출력하는게 정해인데 random.sample()을 이용한 트릭으로 풀었다 ㅋㅋ

 

정석대로 풀면 다음과 같다.

import itertools

L = sorted(map(int, open(0)))
for i, j in itertools.combinations(range(9), 2):
    s = sum(L) - L[i] - L[j]
    if s == 100:
        for k in range(9):
            if k == i or k == j:
                continue
            print(L[k])
        break

 

당연한 얘기지만 배열 입력이 한 라인에 주어지는 경우에는 적용할 수 없다.