백준1256번: 사전

dp와 조합론으로 푸는 문제이다.

import sys
from math import comb
input=sys.stdin.readline
n,m,k=map(int,input().split())
result=[]
if comb(n+m,n)<k:
    print(-1)

else:
    while n>0 and m>0:
        cb=comb(n+m+1,n-1)
        if cb>=k:
            result.append('a')
            n=n-1

        else:
            result.append('z')
            m=m-1
            k=-cb

    result.extend(['a']*n)
    result.extend(['z']*m)

    print(''.join(result))

처음에 풀었을때는  comb(n+m, n)은 (n + m)개 중에 n개의 'a'를 선택하는 경우의 수를 이용해, k보다 작으면 문자열을 만들 수 없다는 의미로 -1을 출력하는 조건문을 바로 위에다 걸어두고, cb라는 변수에 comb(n+m-1, n-1)함수를 넣어서 현재 위치에 'a'를 넣는 경우의 수라고 생각을해 조건을 걸어주었다. cb >= k이면 'a'를 추가하고, 그렇지 않으면 'z'를 추가했다. z'를 추가하는 경우 k 값을 줄여주었다. 'z'를 추가한 경우 나머지 부분에서 k번째 문자열을 찾기 위함이였다. 모든 루프가 끝난 후 남은 'a'와 'z'를 결과 리스트에 추가하였다.

 

그러나 이렇게 하면 cb는 'a'를 선택했을 때 남은 문자열의 경우의 수이기 때문에,이 값을 그대로 빼면 k값이 실제로 원하는 위치에서 벗어날 수 있다.

import math

def find_kth_string(N, M, K):
    total_strings = math.comb(N + M, N)  # 전체 문자열의 개수는 조합(N+M, N)과 같다.

    if K > total_strings:
        return -1  # K가 문자열의 개수보다 크면 -1을 반환

    result = []
    while N > 0 and M > 0:
        combinations = math.comb(N + M - 1, N - 1)  # 현재 위치에서의 조합을 계산
        if K <= combinations:
            result.append('a')
            N -= 1
        else:
            result.append('z')
            M -= 1
            K -= combinations

    result.extend(['a'] * N)
    result.extend(['z'] * M)

    return ''.join(result)

# 입력 예제
N, M, K = map(int, input().split())

# 결과 출력
answer = find_kth_string(N, M, K)
print(answer)

 

이게 정답 코드로, K가 combinations보다 작거나 같으면 'a'를 선택하게하고, 그렇지 않으면 'z'를 선택한다음,  K 값을 combinations만큼 감소시켜 남은 문자열에서 K번째를 찾게하였다.