[Codechef] Logan and DIGIT IMMUNE numbers

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : Logan and DIGIT IMMUNE numbers
Source            : Codechef
Category          : Dynamic Programing
Algorithm         : Digit DP
Verdict           : Accepted 

#include "bits/stdc++.h"
 
using namespace std;
 
#define ll              long long
 
inline int setbit(int mask, int pos)        { return mask |= (1 << pos); }
inline int resetbit(int mask, int pos)      { return mask &= ~(1 << pos); }
inline int togglebit(int mask, int pos)     { return mask ^= (1 << pos); }
inline bool checkbit(int mask, int pos)     { return (bool)(mask & (1 << pos)); } 
 
ll dp[20][1100][3][5][7][9]; // dp[ position ][ mask ][ % 3 ][ % 5 ][ % 7 ][ % 9 ]
bool memoize[20][1100][3][5][7][9];
int a[20];
 
inline ll compute(int mask, int three, int five, int seven, int nine)
{
      for (int i = 0; i <= 8; i+=2) if (checkbit(mask, i)) return 0;
      if (checkbit(mask, 1)) return 0;
      if (checkbit(mask, 3) && three == 0) return 0;
      if (checkbit(mask, 5) && five == 0) return 0;
      if (checkbit(mask, 7) && seven == 0) return 0;
      if (checkbit(mask, 9) && nine == 0) return 0;
      return 1;
}
 
inline ll solve(int pos, int mask = 0, int three = 0, int five = 0, int seven = 0, int nine = 0, bool flag = 1, bool take = 0)
{
      if (pos < 0) return compute(mask, three, five, seven, nine);
      ll &ret = dp[pos][mask][three][five][seven][nine];
      bool &mem = memoize[pos][mask][three][five][seven][nine];
      if (!flag && take && mem) return ret;
      int loop = 9;
      if (flag) loop = a[pos];
      ll sum = 0;
      for (int i = 0; i <= loop; i++)
      {
            if (!take && i == 0) sum += solve(pos-1, 0, 0, 0, 0, 0, 0, 0);
            else sum += solve(pos-1, setbit(mask, i), (three*10+i)%3, (five*10+i)%5, (seven*10+i)%7, (nine*10+i)%9, flag & i == a[pos], 1);
      }
      if (!flag && take) mem = 1, ret = sum;
      return sum;
}
 
inline ll digit(ll num)
{
      int len = 0;
      while (num)
      {
            a[len++] = num % 10;
            num /= 10;
      }
      return solve(len-1);
}
 
ll kth(ll a, ll b, ll k)
{
      ll low = a;
      ll high = b;
      ll ans = -1;
      ll precal = digit(a-1);
      while (low <= high)
      {
            ll mid = low + (high - low) / 2;
            ll total = digit(mid) - precal;
            if (total == k)
            {
                  ans = mid;
                  high = mid - 1;
            }
            else if (total > k) high = mid - 1;
            else low = mid + 1;
      }
      return ans;
}
 
int main()
{
      int q;
      cin >> q;
      while (q--)
      {
            ll a, b, k;
            cin >> a >> b >> k;
            cout << kth(a, b, k) << endl;
      }
}