(구름LEVEL) ab를 bba로

문제

https://level.goorm.io/exam/51356/ab%EB%A5%BC-bba%EB%A1%9C/quiz/1

풀이

처음에는 스택 두 개를 두고 푸는 방법을 생각했다. 특정 위치의 문자가 'b'일 경우 그 전의 문자가 무엇이었는지 스택의 top을 통해 알아내는 방법으로 다음과 같이 구현하였다.

코드 1

import java.io.*;
import java.util.*;

class Main {
	
	static int value = 1000000007;
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		String input = br.readLine();
		
		Stack<Character> s1 = new Stack<>();
		Stack<Character> s2 = new Stack<>();
		long count = 0;
		
		for (int i = 0; i < input.length(); i++) {
			char c = input.charAt(i);
			
			if (c == 'a' || (s1.isEmpty() || s1.peek() == 'b')) {
				s1.push(c);
			} else {	
					count = (count + 1) % value;
					s1.pop();
					s2.push('a');
					s2.push('b');
					s2.push('b');
					
					while (!s2.isEmpty()) {
						char tmp = s2.pop();
						
						if (tmp == 'a' || (s1.isEmpty() || s1.peek() == 'b')) {
							s1.push(tmp);
						} else {
							count = (count + 1) % value;
							s1.pop();
							s2.push('a');
							s2.push('b');
							s2.push('b');
						}
					}
			}
		}
		
		System.out.println(count);
	}
}

답을 구하는데는 문제가 없었는데 런타임 에러가 났다. 혹시 놓친 부분이 있었나 살펴봤지만, 그 문제가 아니었다.
문자열의 길이는 20만인데, ab를 bba로 바꾸는 작업을 진행하면서 스택에 담긴 원소의 개수는 급격히 증가하여 결국 스택은 터지게 되는 것이었다.

이렇게 문제가 발생한다는 것은, dp 문제일 가능성이 높다는 것이다. 가만히 보면 ab를 bba로 변환했을 때 문자열의 길이와 a의 위치가 달라질 뿐 a의 개수는 유지된다는 것을 알 수 있었다.

ex)
1. aaab -> aabba
2. aabba -> abbaba
3. abbaba -> bbababa
4. bbababa -> bbabbbaa
5. bbabbbaa -> bbbbabbaa
6. bbbbabbaa -> bbbbbbabaa
7. bbbbbbabaa -> bbbbbbbbaaa


aaab와 bbbbbbbbaaa의 a는 각각 3개로 동일하다.

그렇다면, b 앞에 있는 a의 개수에 대한 규칙이 있다고 생각하였고, 다음과 같은 규칙을 얻을 수 있었다.

규칙
ab -> 1
aab -> 3
aaab -> 7
aaaab -> 15
...

dp[i] = dp[i - 1] * 2 + 1 라는 사실을 알 수 있었다.
(단, 문제에서 제시해준 모듈러 연산은 고려해줘야 한다.)

이러한 성질을 이용하여 코드를 다시 작성하였다.

코드 2

import java.io.*;
class Main {
	
	final static int value = 1000000007;
	
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int answer = 0;
		int[] dp = new int[1000001];
		
		for (int i = 1; i < dp.length; i++) {
			dp[i] = (dp[i - 1] * 2 + 1) % value;
		}

		String input = br.readLine();
		
		int count = 0;
		for (int i = 0; i < input.length(); i++) {
			char c = input.charAt(i);
			
			if (c == 'a') {
				count++;
			} else {
				answer = (answer + dp[count]) % value;
			}
		}
		
		System.out.println(answer);
	}
}