[Python] List Comprehension

코딩테스트 문제를 풀다 보면 문제를 입력받을 때, 이차원 배열을 받을 때가 많다.

for문을 통해서 배열을 입력받을 때 빈 리스트에 한 줄씩 append 해줄 수 있지만,

list comprehension을 이용하면 쉽게 한줄로 입력받을 수 있다.

개인적으로 입력을 받거나 행렬을 생성할 때 간단하게 리스트를 표현할 수 있기 때문에 특히 많이 사용하게 되는 것 같다.

 

• 표현식 + for문 형식
  result = [표현식 for 변수 in 리스트]

# 예시
# 1차원 배열
# 숫자를 순차적으로 입력받을 때
result = []
for i in range(n):
    result.append(int(input()))

# list comprehension
result = [int(input()) for i in range(n)] 


# 2차원 배열
# 1 2 3
# 4 5 6
# 위의 예시처럼 배열을 차례로 입력받을 때
result = []
for i in range(n):
    result.append(list(map(int, input().split())))

# list comprehension
result = [list(map(int, input().split())) for i in range(n)] 


# 0으로 초기화한 2차원 배열 생성할 때
result = []
for i in range(n):
    result.append([0]*n)
   
# list comprehension
result = [[0]*n for i in range(n)]

 

• 표현식 + for문 + 조건문 형식
  result = [표현식 for 변수 in 리스트 조건문]

# 예시
# 짝수 리스트를 만들 때 사용될 수 있다. 
result = []
for i in range(n):
    if i % 2 == 0:
        result.append(i)

# list comprehension
result = [i for i in range(n) if i % 2 == 0]

# [0,2,4,6...]

 

• True 표현식 + if 조건문 else + False 표현식 + for문 형식 
  result = [True 표현식 if 조건문 else False 표현식 for 변수 in 리스트]
  else까지 써야 하는 경우라면 조건문이 for문보다 선행한다.

# 예시
# 홀수일 땐 True, 짝수일 땐 False인 리스트를 만들 수 있다. 
result = []
for i in range(n):
    if i % 2 == 0:
        result.append(False)
    else:
        result.append(True)

# list comprehension
result = [False if i % 2 == 0 else True for i in range(n)]

# [False,True,False,True,False,True...]

 

• 중첩 for문 형식
  result = [표현식 for 변수1 in 리스트 1 for 변수 2 in 리스트 2...]

# 예시
# 두 개의 리스트가 주어졌을 때, 서로 쌍을 만들 때 사용될 수 있다.
a = [1,2]
b = [9,8,7]

result = []
for i in a:
    for j in b:
        result.append((i,j))

# list comprehension
result = [(i,j) for i in a for j in b] 

# [(1,9), (1,8), (1,7), (2,9), (2,8), (2,7)]