けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

CODE FESTIVAL 2018 qual A C - 半分 (500 点)

少し重たい

問題へのリンク

問題概要

長さ  N の整数列  A_1, A_2, ..., A_N が与えられます。 この数列に以下の操作をちょうど  K 回施します。

  • 添字  i を一つ選んで、 A_i を 2 で割る (小数点以下切り捨て)

 K 回の操作のあとの数列としてありうるものの個数を 109+7 で割ったあまりを求めよ。

制約

  •  1 \le N \le 50
  •  0 \le K \le 10^{9}

考えたこと

 N が小さい...
そして  K が大きい...

素直な DP を考えると、

  • dp[ i ][ k ] := 最初の i 個について k 回の操作で得られる数列の個数

とすればよさそう。だが単純にやると以下のようなループになって  O(NK^{2}) かかってしまってとんでもない。

for (int i = 0; i < N; ++i) {
  for (int k = 0; k <= K; ++k) {
    for (int k2 = 0; k + k2 <= K; ++k2) {
      if (なんらかの条件)
        dp[i][k][k+k2] += dp[i][k][k];
    }
  }
}

まず、要素を 1 つ固定して、その要素を操作によって何種類の数が形作られるのかを考えてみる。たとえば 37 だと

37, 18, 9, 4, 2, 1, 0

の 7 個の数が作られうる。そして 6 回の操作で 0 になることがわかる。0 になったらそれ以上操作をやっても 0 のままであることに注意。

そこで以下のように場合分けできる:

  • 操作後の数列に 0 がある場合: 操作回数が  K 回未満でも 0 に無限に操作を繰り返すことで、数列を変化させることなく操作回数を  K に伸ばせる
  • 操作後の数列に 0 がない場合: 操作回数はちょうど  K 回でなければならない

前者については、操作回数は  K以下であるとしても構わない。そしてダブルカウントが怖いのでそれを防ぐために


一度 0 になった要素にはもう操作してはいけない


というルールを追加する。そうして

  • 操作後の数列に 0 があるような、操作回数が  K 回未満の数え上げ
  • 操作後の数列に 0 がないような、操作回数が  K 回ちょうどの数え上げ

を実施すればよい。さらに考察を進めると、実は 1 要素あたり最大でも 60 回で 0 になってしまうので、50 × 60 = 3000 回以上の操作はもはや考える必要がなかったりする。安全をとって  K > 5000 くらいになると、数列を全部 0 にしてしまってもまだ操作回数が余ってしまう。なので  K \le 10^{9} という制約はもはや空気である。

以後、以下のような dp をする。

  • dp[ i ][ k ][ 1 ] := 数列の最初の i 個について、k 回の操作を実施して、操作後に 0 があるような場合の数
  • dp[ i ][ k ][ 0 ] := 数列の最初の i 個について、k 回の操作を実施して、操作後に 0 がないような場合の数

 k <  5000 (O(N \log{a_i}) くらいまで考えればよくて、1 回の操作で追加できる操作回数も上限 60 回 ( O(\log{a_i})) とかなので、計算時間は  O(N^{2} (\log{a_i})^{2}) になる。

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

const int MOD = 1000000007;
long long dp[55][5100][3];

int main() {
    int N, K;
    cin >> N >> K;
    vector<long long> A(N), num(N, 0);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        while (A[i] > 0) A[i] /= 2, ++num[i];
    }

    dp[0][0][0] = 1;
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j <= 5000; ++j) {
            // 0 -> 0
            for (int k = 0; k < num[i]; ++k) {
                dp[i+1][j+k][0] += dp[i][j][0];
                dp[i+1][j+k][0] %= MOD;
            }
                
            // 0 -> 1
            dp[i+1][j+num[i]][1] += dp[i][j][0];
            dp[i+1][j+num[i]][1] %= MOD;
                
            // 1 -> 0
            for (int k = 0; k <= num[i]; ++k) {
                dp[i+1][j+k][1] += dp[i][j][1];
                dp[i+1][j+k][1] %= MOD;
            }
        }
    }

    long long res = 0;
    // 0
    if (K <= 5000) res += dp[N][K][0], res %= MOD;
    
    // 1
    for (int k = 0; k <= min(K, 5000); ++k) res += dp[N][k][1], res %= MOD;
        
    cout << res << endl;
}