문제

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

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

baekjoon

코드

#include <iostream>
#include <vector>

using namespace std;

int main() {
  int n, k;
  cin >> n >> k;
  
  for (int a = 0; a <= n; a++) {
    int b = n - a;
    if (a * b < k) continue;

    vector<int> a_cnt(b + 1, 0);
    for (int i = 0; i < a; i++) {
      int x = min(k, b);
      a_cnt[x]++;
      k -= x;
    }

    for (int i = b; i >= 0; i--) {
      for (int j = 0; j < a_cnt[i]; j++) cout << 'A';
      if (i > 0) cout << 'B';
    }
    return 0;
  }

  cout << -1;
}

풀이

B를 b개 세워두고, A를 a개만큼 적절한 위치에 배치하는 방식이다. a + b는 n개가 되야하고, a * b가 만들 수 있는 AB쌍의 최대 개수다. 만약 a * b < k 라면, 반드시 문자열을 만들 수 있게 된다. a_cnt[i]는 i번째 위치에 들어갈 A의 개수가 된다. 가장 왼쪽에 A를 두면 b개 만큼의 쌍이 생긴다. 그러므로 k와 b 중 작은 값에 대해서 해당 위치에 A를 세우고, 그만큼 k에서 감소시킨다. k가 0이 되더라도 맨 오른쪽에 남은 a만큼 A를 세우면 AB쌍이 생기지 않으므로 괜찮다. 그 후, 각 B에 대해 앞에 올 A의 개수만큼 A를 먼저 출력하고, B를 세운다. 맨 끝 B는 세우지 않는다.