[TIL 210610] 파이썬 List 내포 활용하기 (Leetcode #566)

Q)

m * n 행렬의 형태를 취하는 리스트 mat을 r * c 행렬 형태로 변환하여 출력하고, 정상적인 변환이 불가능할 경우, mat을 그대로 출력하라.

input example : mat = [[1,2],[3,4]], r = 1, c = 4

output = [[1,2,3,4]]

풀이과정)

우선 m * n == r * c 가 성립해야 변환되는 행렬의 빈 칸에 원소가 빠짐없이 들어갈 수 있으므로 조건문을 추가한다.

mat은 2차원 리스트이므로, mat 행렬 안의 원소만을 나열한 1차원 리스트를 만들기 위해 itertools.chain 함수를 활용했다.

from itertools import chain
elements = list(chain(*mat)) # [1,2,3,4]

if r*c == len(elements):  # elements 원소 개수 == m*n
    output = []
    for i in range(r):
        row = [] 			# 다음 i로 넘어가서 row를 초기화 후,
        for j in range(c):
            row.append(elements[i*c+j]) # 직전 i에서 row에 마지막으로 append한 값의 바로 다음 값 인덱스부터 시작할 수 있게 한다.
            output.append(row) 		
    return output			# [[1,2,3,4]]
return mat

더 깔끔한 풀이)

굳이 import를 사용할 필요 없이, 리스트 내포 함수를 활용한다면 elements 리스트를 깔끔하게 만들 수 있다.

elements = [x for row in mat for x in row] # mat안의 row안의 x 반복문

리스트 내포에서 2중 반복문을 바로 생각해내기 어려워 결국 chain을 썼다.

또,

if r * c == len(elements):
    return [elements[i*c : (i + 1)*c] for i in range(r) ] # elements의 값들을 col 개수씩 슬라이싱하는 것을 row 수만큼 반복하고 리스트로 감싸기

위 one-liner처럼 리스트 내포에서 range(r)만큼 인덱스 슬라이싱을 반복했다면 굳이 객체를 여러 개 만들어 append를 반복할 필요도 없었을 것이다.