브론즈 레벨의 문제이지만 나름대로 고민하면서 수의 규칙을 찾아 풀려고 했다. 한층 한층 적어가던 중 어디서 봤던 것 같은 패턴이 눈에 들어오기 시작했다.
바로 파스칼의 삼각형.
처음엔 가우스와 관련된건줄 알고 가우스 키워드로 한참 찾아 헤매다가 파스칼의 삼각형인 것을 알아냈다.
재귀함수
처음에는 재귀함수를 사용해서 풀려고 했다. 어떤 방이 주어졌을 때, 아래층 1호부터 바로 아래 호수까지의 입주자 수를 다 더하고, 또 아래층의 각 방마다 그렇게 하고, 이렇게 0층까지 계속 이어져야 한다면, 재귀적인 방법이 직관적인 것 같다. 답은 바르게 나오는데 시간 초과가 나서 다른 방법으로 풀어야 했다. 아래는 구현한 코드이다.
# 재귀함수를 사용한 방법
_, *a = map(int, open(0).read().split())
def num_inhabitants(k: int, n: int) -> int:
inhabitants = 0
if k == 0:
return n
for i in range(1, n + 1):
inhabitants += num_inhabitants(k - 1, i)
return inhabitants
for i, j in zip(a[::2], a[1::2]):
print(num_inhabitants(i, j))
먼저 num_inhabitants 함수를 정의해준다. 정수 k 와 n(k층 n호)을 받아서 정수값을 리턴하는 함수이다. 아래의 for 반복문에서 k - 1 값을 되먹여서 이 함수를 다시 호출하기 때문에 그 위에 탈출조건을 정의한다. k가 0이 되었을 때 받은 n을 반환하고 종료, 밑에서 함수가 다시 호출되어 k가 음수가 되었을 때 이제까지 더한 inhabitants를 반환하고 종료한다.
그리고 각 케이스 당 print(num_inhabitants(...)) 함수를 호출하면 k층 n호의 입주자 수를 출력한다.
앞서 말했던 것처럼 약간 느린 방법이다. 시간 초과가 나서 다른 방법을 강구해야 했다.
파스칼의 삼각형
숫자들을 종이에 적다 보니 패턴이 눈에 보이기 시작했다.
파스칼의 삼각형이다.
(출처: 나무위키)
0층에서는 1씩 증가하는데 1층부터는 증가율이 증가하기 시작한다. 각 층의 증가율의 증가율은 바로 아래층의 증가율이다.
구현
파스칼의 삼각형을 구현하기 위해 먼저 math 라이브러리의 factorial 함수를 임포트했다. (나중에 숏코딩에서 볼 수 있지만 comb()함수를 사용해 바로 구할 수도 있었다.)
from math import factorial
_, *a = map(int, open(0).read().split())입력값의 첫번째 값을 제외하고 a에 저장한 뒤 zip함수와 리스트 슬라이싱을 사용해 묶어서 바로 for 반복문에 투입시킨다.
for k, n in zip(a[::2], a[1::2]):
k += 1
ls = []k를 1 증가시키는 이유는 예를 들어 k 입력값이 5일 때(5층) 0층부터 시작하므로 6차(k + 1)의 이항계수(1 6 15 20 15 6 1)를 사용할 것이기 때문이다. 리스트 ls를 초기화시키고 for 반복문 안에 다른 for 반복문을 사용하여 ls에 이항계수를 append한다.
for i in range(0, k + 1):
ls.append(factorial(k)/(factorial(i)*factorial(k-i)))
여기서 math.comb함수로 ls.append(comb(k, i))로 간단하게 작성할 수도 있다. k차 이항계수의 i번째 값을 뜻한다.
가장 바깥의 반복문 안에 for 반복문을 추가하여 리스트에 덧셈을 해주고, 그 첫번째 값이 답이 될 수 있도록 했다.
for _ in range(n - 1):
for i, _ in enumerate(ls):
try:
ls[i] += ls[i + 1]
except IndexError:
pass
print(int(ls[0]))
ls의 각 원소에 다음 원소를 더해주고 마지막 값은 증가율이 0이므로 IndexError에서 pass. n번 반복한 뒤 결과값을 출력한다.
완성된 코드:
from math import factorial
_, *a = map(int, open(0).read().split())
for k, n in zip(a[::2], a[1::2]):
k += 1
ls = []
for i in range(0, k + 1):
ls.append(factorial(k)/(factorial(i)*factorial(k-i)))
for _ in range(n - 1):
for i, _ in enumerate(ls):
try:
ls[i] += ls[i + 1]
except IndexError:
pass
print(int(ls[0]))
숏코딩 분석
고수들은 어떻게 했을까?
숏코딩 또한 파스칼의 삼각형을 사용한다.
# Originally:
import math
i=input
for n in[int]*int(i()):k=n(i());print(math.comb(k+n(i()),k+1))
# Equivalently:
import math
T = int(input())
for n in [int] * T: # [int]는 뭐지? -> ##
k = n(input()) ## n이 이 줄의 input()값을 정수형으로 변환해준다.
# comb 함수로 factorial을 대신할 수 있었다.
print(math.comb(k + n(input()), k + 1)) # 왜 comb(k + n, k + 1)인가?
# factorial(k + n)/(factorial(k + 1)*factorial(n - 1))
# 여기서 k + n 은 파스칼의 삼각형을 확장시켜서 n 차 더 올라가는 것이고
# (예를 들어 4층 4호 -> 8차, 4층 1호 -> 5차)
# k + 1 은 k + n 차의 이항계수에서 k + 1번째의 숫자를 선택하게 한다.
# n - 1 이어도 된다. (식에서 볼 수 있듯) 같은 값을 준다.
파스칼의 삼각형 n번째 줄 k번째 원소의 값을 comb(n - 1, k - 1)로 표현할 수 있다.
파스칼의 삼각형을 가져다가 반시계방향으로 약간 회전시켜보자.
0층 1 2 3 4 5 6 7 ...
1층 1 3 6 10 15 21 28 ...
2층 1 4 10 20 35 56 84 ...
3층 1 5 15 35 70 126 210 ...
이런 식으로 진행됨을 볼 수 있다.
예를 들어 3층 3호의 입주민 수가 궁금하다고 하자. 첫번째 줄을 1, 1, 1, ... 라고 했을 때, 1, 2, 3, 4, ... 가 0층이므로, 대각선으로 다섯번째 줄이 3층에 해당한다. 3호는 그 줄에서 3번째 숫자인 15인데, 이는 6차 이항계수의 3번째 숫자임을 볼 수 있다(1, 6, 15, 20, 15, 6, 1 에서 3번째 숫자). 일반화하면, k층 n호의 값은 파스칼의 삼각형 해당 원소의 값을 구하는 함수인 comb(k + n, n - 1) == factorial(k + n)/(factorial(k + 1)*factorial(n - 1)) == comb(k + n, k + 1)로 구할 수 있다.
Note
뿌듯하다. 8전 9기.
규칙을 발견해서 풀었다.
손으로 직접 쓴게 도움이 되었다.
파스칼의 삼각형을 이해하는데 도움을 주신 홍엽님과
같이 코드리뷰를 해주신 태양님 감사드립니다.
멀티캠퍼스 파이썬 풀스택 취업캠프 2기 화이팅!
