(1040) 정수

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

단계별로 풀어보기

50단계(동적 계획법 4) 10번째

 

 

 

Digit DP에 대해 배우는 문제.

DP[a][b][c]를 다음과 같이 정의한다.

a: 현재 탐색중인 자리수 인덱스

b: 사용한 숫자를 나타내는 비트마스크. x를 사용했다면 2^x 비트가 켜져 있다.

c: 지금까지 만들어진 숫자가 N보다 큰지 여부

값: 다음에 선택 가능한 최적, 최소의 한자리 숫자. 불가능하다면 10

 

DP[a][b][c]에서 숫자를 선정하는 과정은 다음과 같다. c가 0이라면 N보다 커져야 하므로 현재 자리수에 해당하는 숫자부터, 1이라면 이미 N보다 크므로 0부터 9까지 탐색한다. 오름차순으로 탐색하므로 처음 return되는 10이 아닌 수가 정답의 일부가 된다.

다음 숫자를 탐색해 나가다가 마지막 자리에 도달했다면 b를 검사해 K종류의 숫자를 사용했을 때만 성공을, 아니면 실패를 반환한다.

DP 정의에 의해 갱신이 끝난 뒤 기저 상태부터 시작해 다음 숫자를 계속해서 따라가면 정답을 찾을 수 있다.

 

Digit DP를 최적화하는 방법이 몇 가지 있는데, DP배열에서 c=1인 경우는 c=0인 경우에 비해 매우 적기 때문에 DP[a][b]로 c=0인 경우만 저장하고 c=1인 경우에는 메모이제이션 대신 다시 계산하도록 하면 시간을 미세하게 희생해서 메모리를 절반으로 줄일 수 있다. 지금까지 앞자리에 0만 왔는지 여부도 필요한 경우가 있는데(이 문제의 c도 써서 DP[a][b][c][d]인 경우도 있다), 이런 것도 저장 없이 메모리를 줄일 수 있다.

그리고 코드 특성 상 캐시 미스가 자주 일어나기 때문에 벡터 대신 배열로 선언하거나 계산식을 통해 1차원으로 줄이는 방법도 있다.

 

 

 

#include <iostream>
#include <vector>
using namespace std;

int main() 
{
  string n;
  int k;
  
  cin >> n >> k;
  
  vector dp(n.length()+1, vector<vector<int>> (1024, vector<int> (2, -1)));
  
  auto func= [&](auto self, int digit, int state, int tight) -> int {
    if(digit==n.length())  return (__builtin_popcount(state)==k ? 1 : 10);
    
    int &ret= dp[digit][state][tight];
    
    if(ret!=-1)  return ret;
    
    int limit= (tight ? n[digit]-'0' : 0);
    
    for(int i=limit;i<=9;i++){
      int nowtight= tight && (i==limit);
      int nowstate= state|(1<<i);
      
      if(__builtin_popcount(nowstate)<=k){
        int res= self(self, digit+1, nowstate, nowtight);
        
        if(res!=10){
          return ret= i;
        }
      }
    }
    
    return ret= 10;
  };
  
  auto trace= [&]() -> long {
    long res= 0;
    
    for(int digit=0,state=0,tight=1;digit<n.length();digit++){
      int pick= dp[digit][state][tight];
      res= res*10+pick;
      
      int limit= (tight ? n[digit]-'0' : 0);
      tight= tight && (pick==limit);
      state= (state | (1<<pick));
    }
    
    return res;
  };
  
  while(true){
    int ans= func(func, 0, 0, 1);
    
    if(ans!=10){
      cout << trace();
      break;
    }
    
    for(int i=0;i<n.length();i++)  n[i]= '0';
    n= "1"+n;
  }
  
  return 0;
}