CFR 431 D. Random Task ( Binary Search, Digit DP )

Problem - 431D - Codeforces

題意：
給 M, K。求一個 N，滿足 1 ≤ N ≤ 1e18 且在二進制下，[ N + 1, 2 * N ] 中恰有 M 個數恰有 K 個 1。

資料規模：
The first line contains two space-separated integers, m and k (0 ≤ m ≤ 1e18; 1 ≤ k ≤ 64).

解法：
直覺會先想到數位統計 DP。但這個 DP 只能快速計算一個特定的 N 下有幾個數有 K 個 1。這提示我們尋找某種單調性，因此首先考慮是否有二分性質。
考慮 N 對應的 [ N + 1, 2 * N ] 和 N + 1 對應的 [ N + 2, 2 * ( N + 1 ) ]，不難發現，N + 1 和 2 * ( N + 1 ) 在二進制下的意義是等價的，因為它們的 1 的個數一定一樣。如此一來，由於額外獲得了 2 * ( N + 1 ) - 1，所以 N + 1 的 DP 值一定不會比 N 少。
dp[ i ][ j ][ k ][ l ] : 在二進制下，已從最高位考慮前 i 個數字，確定比 2 * N 小，確定比 N + 1 大，有 l 個 1，此時的方案數。

時間 / 空間複雜度：
O( K * lg^2 N )

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

ll dp[ 64 + 1 ][ 2 ][ 2 ][ 64 + 1 ];

ll calc( ll val, const int gol ){
  vector< int > hi, lo;
  for( ll i = val * 2; i; i >>= 1 ){
    hi.emplace_back( i & 1 );
  }
  for( ll i = val + 1; i; i >>= 1 ){
    lo.emplace_back( i & 1 );
  }
  while( hi.size() > lo.size() ) lo.emplace_back( 0 );
  reverse( hi.begin(), hi.end() );
  reverse( lo.begin(), lo.end() );
  memset( dp, 0, sizeof( dp ) );
  dp[ 0 ][ 0 ][ 0 ][ 0 ] = 1;
  for( int i = 0; i < hi.size(); ++i ){
    for( int j = 0; j < 2; ++j ){ // less than hi
      for( int k = 0; k < 2; ++k ){ // more than lo
        for( int l = 0; l < 64; ++l ){
          for( int d = k ? 0 : lo[ i ]; d <= ( j ? 1 : hi[ i ] ); ++d ){
            dp[ i + 1 ][ j or d < hi[ i ] ][ k or d > lo[ i ] ][ l + d ] += dp[ i ][ j ][ k ][ l ];
          }
        }
      }
    }
  }
  ll res = 0;
  for( int i = 0; i < 2; ++i ){
    for( int j = 0; j < 2; ++j ){
      res += dp[ hi.size() ][ i ][ j ][ gol ];
    }
  }
  return res;
}

signed main(){
  ios::sync_with_stdio( 0 );
  ll M; int K; cin >> M >> K;
  ll ans = -1;
  for( ll lb = 1, rb = 1e18; lb <= rb; ){
    ll mid = lb + ( rb - lb ) / 2;
    if( calc( mid, K ) >= M ){
      ans = mid;
      rb = mid - 1;
    } else{
      lb = mid + 1;
    }
  }
  cout << ans << endl;
  return 0;
}