백준 1256번 : 사전 in Python
이 문제는 N개의 "a"와 M개의 "z"로 이루어진 문자열만 알파벳 순서대로 수록한 사전에서 K번째 문자열을 찾는 문제이다.
이 문제를 해결하기 앞서, N개의 "a"와 M개의 "z"로 이루어진 서로 다른 문자열의 개수를 생각해 보자. 만약 N개의 "a"와 M개의 "z"가 모두 서로 다른 문자였다면 N+M개의 문자를 이용해 만들 수 있는 문자열의 개수는 (N+M)!일 것이다. 그런데 N개의 "a"와 M개의 "z"는 같은 문자이므로 이들 사이의 순서가 바뀌어도 같은 경우로 간주할 수 있고 따라서 이 경우의 수를 제외해야 한다. N개의 문자를 나열하는 방법의 수는 N!이므로 전체 경우의 수는 (N+M)!/N!M!이다. 이는 N+M개에서 N개를 택하는 조합의 수와 같다.
이를 통해 이 문제는 조합과 관계있음을 알 수 있는데, 특히 알고리즘을 통해 조합을 계산하는 경우 많이 사용하는 방법이 바로 dynamic programming과 조합 공식을 이용하는 것이다. 조합 공식은 n개에서 m개를 택하는 조합의 수를 n-1개에서 m개를 택하는 조합의 수와 n-1개에서 m-1개를 택하는 조합의 수의 합을 통해 구할 수 있다는 것을 알려준다. 이러한 공식이 나오게 된 배경은 다음과 같다. n개에서 m개를 택할 때 n개의 대상 중 맨 처음에 마주하는 대상에 대해 그 대상을 택할지 말지로 경우의 수를 나눌 수 있는데, 그 대상을 택한 경우 나머지 n-1개에 대해 m-1개를 택하면 되고, 택하지 않은 경우 나머지 n-1개에 대해 m개를 택하면 된다. 여기까지 이해했다면 이 문제도 이를 활용해 해결할 수 있다.
먼저 N+M개에서 N개를 택하는 조합의 수를 dynamic programming과 조합 공식을 이용해 구한다. 초깃값으로 임의의 n에 대해 n개에서 0개를 택하는 조합의 수와 n개에서 n개 전부를 택하는 조합의 수를 1로 둔다. 그 다음 i = 2일 때부터 i = N+M일 때까지 순회하고, 그 내부에서 m = 1부터 m = min( i , M+1 )일 때까지 순회하면서 위 조합 공식을 이용해 순차적으로 조합을 계산한다. 그리고 이 결과를 모두 행렬 dp에 저장한다. 즉, dp[i][j]는 i개에서 j개를 택하는 조합의 수를 의미한다.
그런 뒤 이제 사전 내 문자열 중 K번째 문자열을 찾아야 하는데, 우선 K와 dp[N+M][M]를 비교해 K가 더 큰 경우에는 사전에 수록되어 있는 문자열의 개수가 dp[N+M][M]개이므로 K번째 문자열이 존재하지 않는다. 따라서 -1을 출력한다.
그렇지 않은 경우에는 K번째 문자열이 사전 내에 존재한다. 실제로 이 사전을 만든 뒤 K번째 문자열을 찾으려면 매우 큰 메모리가 필요하므로, K번째 문자열을 찾지 않고 직접 만들겠다. 우선 위 조합 공식을 다시 한 번 상기시키면, n개에서 m개를 택하는 조합의 수를 구하기 위해 두 가지 경우로 나눴었다.
이 문제에서는 N+M개의 문자를 가지는 문자열을 만들 때 M개의 "z"를 넣는 방법의 수로 해석하겠다. 그러면 첫 문자가 "z"인 문자열의 개수는 dp[N+M-1][M-1]과 같을 것이고, 첫 문자가 "a"인 문자열의 개수는 dp[N+M-1][M]과 같을 것이다. 이때 사전은 알파벳 순서를 따르므로 첫 문자가 "a"인 문자열이 앞에, "z"인 문자열이 뒤에 온다. 즉, K와 dp[N+M-1][M]을 비교했을 때 K가 더 크고, M이 양수라면 첫 문자가 "a"인 모든 문자열보다 뒤에 있으므로 첫 문자가 "z"일 것이다. 반대로 K가 더 작거나 같거나 M이 0이라면 첫 문자가 "a"일 것이다. 각 경우에 대해 적절한 알파벳을 결과에 붙인다.
그 다음에 첫 문자가 "z"인 경우에는 이제 첫 문자가 "z"인 나머지 dp[N+M-1][M-1]개의 문자열에 대해서만 생각하면 되므로 첫 문자가 "a"인 문자열은 더 이상 고려하지 않는다. 따라서 K에서 첫 문자가 "a"인 문자열의 개수 dp[N+M-1][M]를 빼고, 문자열 내 "z"의 개수인 M을 1만큼 감소시킨다. 첫 문자가 "a"인 경우에는 K와 M을 그대로 둔다. 그리고 두 경우 모두 문자열의 길이인 N+M을 1만큼 감소시킨다.
문자열의 길이가 N+M이 될 때까지 위 알고리즘을 반복하면 원하는 문자열을 얻을 수 있다.
이를 코드로 구현하면 다음과 같다.
N, M, K = map(int, input().split())
# 조합을 dynamic programming을 통해 구하기 위한 행렬
combination = [[0]*(M+1) for _ in range(N+M+1)]
# 초기값 설정
for i in range(1, M+1):
combination[i][i] = 1
combination[i][0] = 1
# 조합 계산
for i in range(2, N+M+1):
for j in range(1, min(i, M+1)):
combination[i][j] = combination[i-1][j] + combination[i-1][j-1]
# K가 사전 내 문자열의 개수보다 많은 경우
if combination[N+M][M] < K:
print(-1)
else:
# K번째 문자열을 저장할 변수
res = ""
# 길이가 N+M인 문자열을 만들기 위해 넣어야 하는 문자의 개수
temp_i = N+M
# 남은 "z"의 개수
temp_j = M
while temp_i > 0:
if K > combination[temp_i-1][temp_j] and temp_j > 0:
K -= combination[temp_i-1][temp_j]
res += "z"
temp_j -= 1
else:
res += "a"
temp_i -= 1
print(res)