[백준 12970][파이썬][그리디] AB

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

12970번: AB

내 기준 상당히 어려웠던 문제. 내가 생각했던 방식과 개선해나가는 과정을 기록해보려고 한다.

 

(A, B) 쌍의 개수 구하기

우선 특정 문자열에서 (A, B) 쌍의 개수는 어떻게 구할 수 있을까?

 

예를 들어 ABBAABBB 라는 문자열에서 (A, B) 쌍의 개수를 구해보자.

 

모든 A에 대해서 각 A 뒤에 있는 B의 개수를 구해서 더하거나 모든 B에 대해서 각 B 앞에 있는 A의 개수를 구해서 더하면 된다.

 

첫 번째 A에 대해서 뒤에 있는 B의 개수는 5, 두 번째와 세 번째 A에 대해서 뒤에 있는 B의 개수는 3이다.

 

즉 5+3*2=11, (A, B) 쌍의 개수는 11개가 된다.

 

A와 B의 개수가 주어질 때 (A, B) 쌍이 가장 많은 문자열

만약 A와 B의 개수가 주어질 때 (A, B) 쌍 가장 많은 문자열은 어떻게 구할 수 있을까?

 

모든 A를 연속하여 붙이고 그 뒤에 모든 B를 붙이면 된다.

 

예를 들어 A의 개수가 4이고 B의 개수가 5라면 (A, B) 쌍의 개수가 최대가 되는 문자열은 무조건 AAAABBBBB이다.

 

조금만 생각해봐도 알 수가 있다.

 

A가 연속한다면 모든 B에 대해서 쌍을 생성할 수 있는데 연속하지 않다면 쌍을 생성할 수 없는 B가 생기기 때문이다.

 

그리디 알고리즘 적용

이제 그리디 알고리즘을 적용해보자.

 

처음에 모두 B로 채워진 문자열을 생성하고 여기에 한 단어씩 A로 변경해가면 된다.

 

그 전에 문자열을 만들 수 없는 경우를 판단해야 한다.

 

위에서 A와 B의 개수가 주어질 때 (A, B) 쌍이 가장 많은 문자열을 구하는 방법에 대해서 알아보았다.

 

만약 문자열의 길이가 6이라면 0*6, 1*5, 2*4, 3*5 ... 이런 식으로 가장 큰 (A, B)쌍의 개수를 구해서 만약 이 값이 k보다 작다면 절대 문자열을 만들지 못한다.

 

import sys
input=sys.stdin.readline
n, k = map(int, input().split())
s=['B']*n
cnt=0 # (A, B) 쌍의 수
# 불가능한 경우 골라내기
max_cnt=0 # 가장 큰 (A, B)쌍의 개수 
for i in range(n//2+1):
    max_cnt=max(max_cnt, i*(n-i))
if max_cnt<k: # 못 만드는 경우
    print(-1)
    exit()

처음에는 두 가지 연산을 생각했었다.

1. 문자열의 앞 부분을 A로 바꾼다.
2. 문자열의 뒷 부분을 A로 바꾸고 한 칸씩 앞으로 땡겨온다.

1번 연산은 다음과 같다.

 

처음에 BBBBB라는 문자열이 있을 때 1번 연산을 적용하면 ABBBB로 바뀌고 한 번 더 1번 연산을 적용하면 AABBB로 바뀐다.

 

A와 B가 연속되어 있을 때 (A, B) 쌍의 개수가 최대가 되므로 k보다 처음으로 커질 때까지 1번 연산을 반복하고 그 다음 2번 연산을 적용하는 방식이다.

 

2번 연산의 의미는 다음과 같다.

 

AABBBB라는 문자열이 있다고 하자. 여기서 2번 연산을 적용하면 AABBBA가 되고 (A, B) 쌍의 개수는 8에서 6으로 감소한다.

 

그리고 한 칸을 땡겨온다면 AABBAB가 될 테고 1이 증가한다.

 

즉 문자열 s에 2번 연산을 적용할 때 처음 마지막 B를 A로 바꿀 때는 문자열 s에 존재하는 A의 개수만큼 (A, B) 쌍의 개수가 감소한다.

 

그리고 한 칸씩 땡겨올 때는 (A, B) 쌍의 개수가 1씩 증가한다.

 

n=7이고 k=11이고 (A, B) 쌍의 개수는 cnt라고 예시를 들어보자.

 

처음에는 BBBBBBB 문자열을 만든다. cnt는 0이다.

 

cnt < k 이기 때문에 1번 연산을 적용한다.

 

ABBBBBB가 되었고 cnt는 6이다.

 

cnt < k 이기 때문에 1번 연산을 적용한다.

 

AABBBBB가 되었고 cnt는 10이다.

 

cnt < k 이기 때문에 1번 연산을 적용한다.

 

AAABBBB가 되었고 cnt는 12이다.

 

cnt > k이기 때문에 가장 마지막 1번 연산을 취소한다.

 

AABBBBB가 되었고 cnt는 10이다.

 

이 때부터 2번 연산을 적용한다.

 

마지막 B를 A로 바꾼다.

 

AABBBBA가 되었고 cnt는 8이다.

 

총 3칸을 땡겨온다.

 

AABABBB가 되었고 cnt는 11이다.

 

이러한 과정을 통해 문자열을 구할 수 있다.

 

그런데 구현이 상당히 까다로운 방식이다. cnt가 처음으로 k를 넘었을 때 그 연산을 취소하고 2번 연산을 수행하는 등 구현이 쉽지 않다.

 

그런데 사실 1번 연산이 아닌 2번 연산만 사용해도 답을 구할 수 있다.

 

2번 연산을 적용해서 계속 앞으로 땡겨오다보면 언젠가는 A가 가장 처음에 도달하게 되고 이는 결국 1번 연산을 적용한 것과 동일하다.

 

물론 1번 연산도 함께 사용하는 것이 시간을 단축할 수 있겠지만 구현의 편의상 2번 연산만 사용하도록 한다.

 

2번 연산의 경우 마지막 B를 A로 바꿀 때 cnt가 줄어들고 한 칸씩 땡겨 올 때마다 cnt가 1씩 증가한다.

 

즉 2번 연산만 사용하면 모든 cnt를 탐색할 수 있다.

BBBBB에서 cnt는 0이고 2번 연산을 계속 적용해보자,

 

문자열	cnt
BBBBB	0
BBBBA	0
BBBAB	1
BBABB	2
BABBB	3
ABBBB	4
ABBBA	3
ABBAB	4
ABABB	5
AABBB	6

보다싶이 모든 경우를 탐색할 수 있다.

 

그런데 만약 A의 개수가 과반수가 넘어가면 어떨까?

 

AAAAAAAABBBB 이러한 문자열에서 2번 연산을 계속 적용해보면 어떨까? 처음에는 cnt가 32이다.

 

 

문자열	cnt
AAAAAAAABBBB	32
AAAAAAAABBBA	24
AAAAAAAABBAB	25
AAAAAAAABABB	26
AAAAAAAAABBB	27
AAAAAAAAABBA	19

28에서 31까지 나오지 않는 모습을 볼 수 있다. 그러면 이 수에 대해서는 탐색을 못하니 모든 경우를 탐색하는게 아닐까?

 

그런데 모두 B로 이루어진 문자열에서 2번 연산만 계속 적용하다보면 A가 과반수가 넘어가는 일 자체가 발생하지 않는다.

 

A와 B의 개수가 주어질 때 (A, B) 쌍이 가장 많은 문자열을 구하는 방법을 다시 생각해보자.

 

길이가 8인 문자열에서 (A, B)쌍이 가장 많은 문자열은 무엇일까?

 

당연히 AAAABBBB 처럼 A와 B의 개수가 가장 비슷한 문자열이다.

 

AAAAABBB 처럼 A의 개수가 과반수가 넘어가는 순간 (A, B) 쌍의 개수가 줄어든다.

 

즉 A의 개수가 과반수가 되기 전까지는 완전 탐색이 맞고 A의 개수가 과반수가 될 때까지 cnt가 k가 되지 못했다면 애초에 만들 수 없는 문자열일테고 가장 처음에 만들 수 없는 문자열은 걸러냈기 때문에 이러한 상황 자체가 발생하지 않는다.

 

<최종 코드>

# 백준 12970
# 그리디
import sys
input=sys.stdin.readline
n, k = map(int, input().split())
s=['B']*n
cnt=0 # (A, B) 쌍의 수
# 불가능한 경우 골라내기
max_cnt=0 # 가장 큰 (A, B)쌍의 개수 
for i in range(n//2+1):
    max_cnt=max(max_cnt, i*(n-i))
if max_cnt<k: # 못 만드는 경우
    print(-1)
    exit()
while cnt!=k:
    idx=n-1 # 뒤에서부터 A 옮기기
    cnt-=s.count('A') # 가장 마지막 B를 A로 바꾸면 앞 A 개수만큼 (A, B)쌍의 수가 줄어듦
    s[idx]='A' # 가장 뒤 B를 A로 변경
    while idx>0 and s[idx-1]=='B'and cnt!=k:
        # A를 한 칸 앞으로 땡겨옴
        s[idx]='B' 
        idx-=1
        s[idx]='A'
        cnt+=1
print(''.join(s))