세 개의 숫자 i, j, k와 또 하나의 숫자 n이 주어져 있다고 가정해 보겠습니다. 이때 i + j + k가 n과 같지 않은 모든 삼중항(i, j, k)의 목록을 찾아야 합니다. 이 문제는 파이썬의 리스트 컴프리헨션(list comprehension) 기법을 사용하면 매우 간결하게 해결할 수 있습니다.
예를 들어 입력이 i = 1, j = 1, k = 2, n = 3일 때 출력은 다음과 같습니다.
[[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 2]]
문제 해결 접근 방법
- x는 0부터 i까지, y는 0부터 j까지, z는 0부터 k까지 범위에서 가능한 모든 조합을 생성합니다.
- 생성된 조합 중 x + y + z가 n과 같지 않은 경우만 골라내어 삼중항 [x, y, z] 형태의 배열(arr)을 만듭니다.
- 리스트 컴프리헨션 안에 조건식(if)을 함께 작성하면 위 과정을 한 줄로 처리할 수 있습니다.
예시 코드
더 잘 이해할 수 있도록 다음 구현 예제를 살펴보겠습니다.
def solve(i, j, k, n): arr = [[x, y, z] for x in range(i+1) for y in range(j+1) for z in range(k+1) if x+y+z != n] return arr i = 1 j = 1 k = 2 n = 3 print(solve(i, j, k, n))
입력
1, 1, 2, 3
출력
[[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 2]]
코드 설명
solve 함수는 세 개의 반복문을 리스트 컴프리헨션 하나로 결합한 형태입니다. range(i+1)처럼 범위 끝에 1을 더해주는 이유는 0부터 i까지의 값을 모두 포함시키기 위해서입니다. 마지막의 if x+y+z != n 조건은 합이 n이 되는 조합(예: 위 예제에서 합이 3인 [1, 0, 2], [0, 1, 2], [1, 1, 1])을 결과에서 제외하므로, 우리가 원하는 삼중항 목록만 깔끔하게 반환됩니다.