yukicoder #362: 門松ナンバー

問題

No.362 門松ナンバー - yukicoder
3桁以上の連続する3桁がすべて門松列である数字を門松ナンバーと呼ぶ。
N番目の門松ナンバーを答えよ。

本番で考えたこと (誤っている方法）

ご丁寧に最大ケースの解を出してくれていて、なおかつ14桁であるという事がわかる。これは桁ごとに計算するDPっぽい。
dp[d][n][f]... d桁で最上桁(d桁目)がnであり、d-1桁目の値＞d桁目の値(f=1)である(f=0の時はこの逆)門松ナンバーの数
とかして計算するとどうだろう、と考える。
Nの値がsum{dp[d][n][f]}となるような臨界点を求めて、それベースで最上桁から順に決めていくとか。
=> 実装しながらこの方針だと駄目だ、と気づく。

何故この方針がダメかというと、例えば、N=sum{dp[3][n][f]}+1であるとき、
解は4桁で最上桁が1であることはわかるが4桁で最上桁が1である門松numberの下3桁は、3桁の最小門松ナンバー(102)ではないから。

じゃあどうしよっか……そもそも2桁以下の扱いとか途中に0のある値の処理が大変だぞ……と思ったところで時間切れ。

解説を読んで

２分探索と組み合わせるとは思いつかなかった。
入力Xを与えた時に、dp[d][n][m]][f]：d桁目がn、d-1桁目がmのd桁の数字で、Xの下d桁よりも大きい(f=1) or 以下である(f=0)門松ナンバー
というdp列を用いて、X以下の門松ナンバーの数を数え上げる。
これができればXを二分探索しながら正解を探せば良い。
2桁以下の値とか途中に0が出るケースは頑張る。

コード

#include <iostream>
#include <climits>
#include <vector>
#include <cstdio>
#include <cstring>
#include <set>
using namespace std;
bool kadomatsu(int x,int y, int z){
   if( x==y || y==z || z==x ) return false;
   if( x<y && y>z ) return true;
   if( x>y && y<z) return true;
   return false;
}
long long func( long long X ){
   long long ret = 0;
   long long dp[16][10][10][2];
   for( int i=0; i<16; i++ ) for( int j=0; j<10; j++ ) for(int k=0 ;k<10; k++ )for(int l=0; l<2;l++){
      dp[i][j][k][l]=0;
   }
   int offset=0;
   for( int i = 0 ; i < 1000; i++ ){
      int x=i/100;
      int y=(i/10)%10;
      int z=i%10;
      int flag = 0;
      if( X%1000 < i ) flag =1;
      if( x==y || y==z || x==z ) continue;
      if( x<y&&y>z ){
         if( x==0 ) offset++;
         if( X>=i ){
            ret++;
         }
         dp[3][x][y][flag]++; 
      }
      if( x>y&&y<z ){
         if( x==0 ) offset++;
         if( X>=i ){
            ret++;
         }
      
         dp[3][x][y][flag]++; 
      }
   }
   long long curDigit=1000;
   int di=4;
   while(curDigit<= X){
      int curD = (X/curDigit)%10;
      for( int i = 0; i <= 9 ; i++ ){
        for( int j=0; j<=9; j++ ){
         for( int k=0; k <=9; k++){
            if(kadomatsu(i,j,k)){
               if( i> curD ){
                  dp[di][i][j][1]+=dp[di-1][j][k][0];
                  dp[di][i][j][1]+=dp[di-1][j][k][1];
               }
               else if( i< curD ){
                  dp[di][i][j][0]+=dp[di-1][j][k][0];
                  dp[di][i][j][0]+=dp[di-1][j][k][1];
               }
               else{
                  dp[di][i][j][0]+=dp[di-1][j][k][0];
                  dp[di][i][j][1]+=dp[di-1][j][k][1];
               }
            }
         }
        }
      }
      curDigit*=10;
      di++;
   }
   for( int d = 4; d <di; d++){
      for( int i=1; i<=9; i++ ){
         for( int j=0; j<=9; j++ ){
            ret+=dp[d][i][j][0];
            if( d!=di-1 ){
               ret+=dp[d][i][j][1];
             }
         }
      }
   }
   return ret-offset;

}

int main(){
   int T;
   cin >> T;
   for( int ti=0; ti<T; ti++ ){
      long long K;
      cin >> K;
      long long left=101;
      long long right=4e13;
      while( left+1 <right ){
         long long mid = (left+right)/2;
         if( func(mid) < K ) left=mid;
         else right=mid;
      }
      cout << right<<endl;
   }
   return 0;        
}