C++로 풀어보는 배열 복원 문제: 공백 없는 숫자 문자열에서 배열 경우의 수 구하기

문제 개요

배열 A의 원소들을 차례대로 출력하는 프로그램이 있다고 가정해 봅시다. 그런데 이 프로그램에 작은 실수가 있어 각 원소 사이에 공백이 삽입되지 않았습니다. 이렇게 얻은 하나의 문자열만으로 원래 배열을 다시 복원할 수 있을까요? 단, 배열의 모든 원소는 1부터 k 사이의 값이라는 조건이 주어집니다.

문자열 s와 정수 k가 주어졌을 때, 배열을 복원할 수 있는 서로 다른 방법의 수를 구해야 합니다. 답이 매우 커질 수 있으므로 10⁹ + 7로 나눈 나머지를 반환합니다.

예를 들어 입력이 S = "1318", k = 2000이라면 출력은 8입니다. [1318], [131,8], [13,18], [1,318], [1,3,18], [1,31,8], [13,1,8], [1,3,1,8]처럼 총 8가지의 서로 다른 배열을 만들 수 있기 때문입니다.

풀이 접근 방법

이 문제는 동적 계획법(DP)과 메모이제이션을 활용하면 효율적으로 해결할 수 있습니다. 각 위치에서 숫자를 잘라내는 경우와 이어 붙이는 경우를 모두 고려하되, 이미 계산한 결과는 저장해 두어 중복 연산을 피하는 것이 핵심입니다.

재귀 함수 설계

모듈러 상수 m = 10⁹ + 7을 정의하고, 결과 저장용 맵 dp를 선언합니다.
add(a, b) 함수는 ((a mod m) + (b mod m)) mod m을 반환하여 오버플로를 방지합니다.
help(idx, s, num, k) 함수의 동작은 다음과 같습니다.
- idx가 문자열 길이 이상이면 1을 반환합니다. 문자열 끝까지 성공적으로 분했다는 의미입니다.
- dp에 이미 계산된 값이 있으면 해당 값을 그대로 반환합니다.
- ret을 0으로 초기화합니다.
- num이 1 이상 k 이하이고 현재 문자가 '0'이 아니라면, 지금까지 누적한 숫자를 하나의 원소로 확정하고 help(idx, s, 0, k)를 재귀 호출합니다.
- num × 10 + 현재 자릿수가 k 이하라면, 숫자를 계속 이어 붙이며 help(idx + 1, ...)을 재귀 호출합니다.
- 계산 결과를 dp에 저장한 뒤 반환합니다.

상향식 DP 메인 로직

n := s의 길이로 설정하고, 크기 n + 1의 배열 ans를 선언한 뒤 ans[0] = 1로 초기화합니다.
인덱스 계산을 편하게 하기 위해 s 앞에 공백 한 칸을 붙입니다.
k를 문자열 ks로 변환하여 자릿수 비교에 활용합니다.
i를 1부터 n까지 반복하면서, j를 i부터 역방향으로 최대 10자리까지 살펴봅니다.
- temp에 현재 문자를 앞쪽에 추가합니다.
- 현재 문자가 '0'이면 선행 0이 되어 유효하지 않으므로 건너뜁니다.
- temp의 길이가 ks보다 길어지면 반복을 종료합니다. k의 최대 자릿수를 초과했기 때문입니다.
- val := stol(temp)로 변환한 뒤, 1 ≤ val ≤ k를 만족하면 ans[i]에 ans[j - 1]을 더합니다.
최종적으로 ans[n]을 반환합니다.

이 방식은 내부 반복이 최대 10자리까지만 진행되므로 전체 시간 복잡도가 O(n × 10), 즉 사실상 O(n) 수준으로 매우 효율적입니다.

구현 예시

아래 코드를 통해 더 잘 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const int m = 1e9 + 7;
class Solution {
   public:
   unordered_map<int, unordered_map<lli, int> > dp;
   lli add(lli a, lli b){
      return ((a % m) + (b % m)) % m;
   }
   int help(int idx, string& s, lli num, int k){
      if (idx >= s.size())
      return 1;
      if (dp.count(idx) && dp[idx].count(num))
      return dp[idx][num];
      int ret = 0;
      if (num >= 1 && num <= k && s[idx] != '0') {
         ret = add(help(idx, s, 0, k), ret);
      }
      if (num * 10 + (s[idx] - '0') <= k) {
         ret = add(help(idx + 1, s, num * 10 + (s[idx] - '0'), k),
         ret);
      }
      return dp[idx][num] = ret;
   }
   int numberOfArrays(string s, int k){
      int n = s.size();
      vector<lli> ans(n + 1);
      ans[0] = 1;
      s = " " + s;
      string ks = to_string(k);
      for (lli i = 1; i <= n; i++) {
         lli cnt = 1;
         string temp = "";
         for (lli j = i; j >= 1 && cnt <= 10; j--, cnt++) {
            temp = s[j] + temp;
            if (s[j] == '0')
               continue;
            if (temp.size() > ks.size())
            break;
            lli val = stol(temp);
            if (val >= 1 && val <= k) {
               ans[i] = add(ans[i], ans[j - 1]);
            }
         }
      }
      return ans[n];
   }
};
main(){
   Solution ob;
   cout << (ob.numberOfArrays("1318",2000));
}

입력

"1318", 2000