백준 12970 파이썬 - AB

https://www.acmicpc.net/problem/12970

12970번: AB

문제

정수 N과 K가 주어졌을 때, 다음 두 조건을 만족하는 문자열 S를 찾는 프로그램을 작성하시오.

• 문자열 S의 길이는 N이고, 'A', 'B'로 이루어져 있다.
• 문자열 S에는 0 ≤ i < j < N 이면서 s[i] == 'A' && s[j] == 'B'를 만족하는 (i, j) 쌍이 K개가 있다.

코드

import sys
n, k = map(int, sys.stdin.readline().split())
a, b = 0, n
while a * b < k and 0 < b:
    a += 1
    b -= 1
if k == 0:
    sys.stdout.write('B' * n)
    exit()
elif b == 0:
    sys.stdout.write('-1')
    exit()
res = k - (a - 1) * b
sys.stdout.write('A' * (a - 1) + 'B' * (b - res) + 'A' + 'B' * res)

풀이

  a를 늘리고, b를 줄이며 a * b < k 인 a, b를 찾는다. 앞에서 부터 a - 1 만큼 A를 넣고, n에서 k - (a - 1) * b를 뺀 인덱스에 A를 넣는다. 나머지는 다 B로 채운다. 수학적으로 생각을 해봐야 하는 그리디 문제이다.