SRM546 Div1 Medium(500), Div2 Hard(1000) FavouriteDigits

FavouriteDigits

MがNより大きいとき、ある桁pについて、pより上位の桁ではNとMが等しく、p桁ではMの数字がNの数字より大きく、pより下位の桁は何でも良い。pより下位の桁は、digit1(d1)＜digit2(d2)ならば、0 0 … 0 d1 d1 … d1 d2 d2 … d2と並べるのが最も小さくなる。pとp桁目の数を変えて、最小のものを返せば良い。

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

class FavouriteDigits{public:
long long findNext( long long N, int digit1, int count1, int digit2, int count2 )
{
    const int M = 16;

    if ( digit1>digit2 )
        swap(digit1,digit2),
        swap(count1,count2);

    vector<int> n(M);
    long long t = N;
    for ( int i=0; t>0; i++, t/=10 )
        n[i] = t%10;

    //  Nが条件を満たすかチェック
    int c1 = count1;
    int c2 = count2;
    int i;
    for ( i=M-1; n[i]==0; i-- )
        ;
    for ( ; i>=0; i-- )
    {
        if ( n[i]==digit1 )  c1--;
        if ( n[i]==digit2 )  c2--;
    }
    if ( c1<=0 && c2<=0 )
        return N;
    
    //  Nより大きく条件を満たす数を探す
    long long ans = 1LL<<62;

    for ( int p=0; p<M; p++ )
    for ( int d=n[p]+1; d<=9; d++ )
    {
        vector<int> t(M);
        for ( int i=M-1; i>p; i-- )
            t[i] = n[i];
        t[p] = d;

        int c1 = count1;
        int c2 = count2;
        int i;
        for ( i=M-1; t[i]==0; i-- )
            ;
        for ( ; i>=p; i-- )
        {
            if ( t[i]==digit1 )  c1--;
            if ( t[i]==digit2 )  c2--;
        }
        c1 = max(c1,0);
        c2 = max(c2,0);

        if ( c1+c2<=i+1 )
        {
            for ( int i=0; i<c2; i++ )
                t[i] = digit2;
            for ( int i=c2; i<c1+c2; i++ )
                t[i] = digit1;

            long long a = 0;
            long long s = 1;
            for ( int i=0; i<M; i++ )
                a += t[i]*s,
                s *= 10;
            ans = min( ans, a );
        }
    }

    return ans;
}};