X보다 작은 수

import sys


def main():
    input = sys.stdin.readline

    n, k = map(int, input().split())
    arr = list(map(int, input().split()))

    smalls = [a for a in arr if a < k]
    print(*smalls)


main()

import sys

def main():
    input = sys.stdin.readline

    n, k = map(int, input().split())
    arr = list(map(int, input().split()))

    smalls = []
    for a in arr:
        if a < k:
            print(a, end=" ")


main()

리스트 컴프리헨션

리스트를 다룰 때, 다음과 같은 연산을 적용하는 경우가 자주 있습니다.

• Filter: 리스트의 원소 중 특정 조건을 만족하는 원소들을 골라냅니다.
• Map: 각 원소에 동일한 연산을 적용합니다.

마치, 집합을 다룰 때 사용하는 조건 제시법(set comprehension)과 유사합니다.

$$ \begin{align*} A &= \{ 1, 2, 3, 4, 5 \} \\ B &= \{ 2x | x \in A, x \le 3 \} = \{ 2, 4, 6 \} \end{align*} $$

위 예시에서, $x \le 3$이 filter, $2x$가 map에 해당합니다.

Python에서는 이와 비슷한 문법을 사용하여 리스트를 다룰 수 있습니다. 이를 리스트 컴프리헨션(list comprehension)이라고 하며, 다음과 같이 사용합니다.

A = [1, 2, 3, 4, 5]
B = [2 * x for x in A if x <= 3]
C = [2 * x for x in A]

B에서 x <= 3이 filter, 2 * x가 map에 해당합니다. C처럼 filter 없이 모든 원소를 사용할 때는 if ... 부분을 생략할 수 있습니다.

멤버 변수를 무시할 수도 있으며, 이때는 관례적으로 _를 사용합니다. 예를 들어, 많은 문제에서 리스트를 입력으로 줄 때, 첫 줄에 원소의 개수 $n$을 주고 그 다음 $n$개의 줄에 걸쳐서 원소를 줍니다.

5
3
1
4
1
5

위 입력으로부터 리스트를 만드는 프로그램은 다음과 같이 작성할 수 있습니다.

n = int(input())
arr = [int(input()) for _ in range(n)]

range(n)는 $n$번 입력을 받기 위해서만 사용하였습니다. arr = [3, 1, 4, 1, 5]가 됩니다.

*-연산자

이터러블의 멤버들을 함수의 인자로 각각 전달하고 싶을 때가 있습니다. Python에서는 *-연산자를 사용하여 이를 구현할 수 있습니다.

def quad(a, b, c, x):
    return ((a * x) + b) * x + c

coeffs = [1, 2, 3]
# 아래 세 줄은 모두 같은 동작을 합니다.
print(quad(*coeffs, -1))
print(quad(*[1, 2, 3], -1))
print(quad(1, 2, 3, -1))

*-연산자는 스플랫 연산자(splat operator)라고도 불립니다.

print 함수는 기본적으로 각 인자를 공백으로 구분하여 출력하기 때문에, 문제의 출력이 이와 같을 때 print(*arr)처럼 사용하면 편리합니다.