CODE FESTIVAL 2018 qual A: C - 半分

https://beta.atcoder.jp/contests/code-festival-2018-quala/tasks/code_festival_2018_quala_c

問題概要

N要素の数列Aが与えられる。数列の要素をひとつ選び2で割る、という操作をちょうどK回行ったあと、残った数列としてありえるものが何通りあるか求めよ。

N <= 50

Ai <= 1018

K <= 109

解法

DPをする。愚直には以下のようなDPとなる。

dp(i, j): i番目までの要素に対してj回操作した結果としてありえる数列の数

問題はKが大きいことで、このDPは遷移も含めるとO(NK2)とかになってしまってまず間に合わない。ここで「ひとつの要素を割り続けると割と早い段階で0になる」「0を割り続ければ状態を変えずに操作回数を空費できる」という2点に注目すると、このDPをもうちょっと効率的にやることができる。

まず Ai <= 1018 より、多くても60回割ればAiは0になる。よって上のDPの操作回数は1要素につき60回、最大で50要素×60=3000回までしか考える必要がない。遷移を考えるときもひとつの要素に対して高々60回まで見ればよい。また操作後の要素にひとつでも0があればそこを割り続けて状態を変えずに余ったKを消費できるので、「操作の結果要素が0になったか」もDPの状態に含めることにする。するとDPは状態数が最大ケースでも 50 * 3000 * 2, 遷移が 60 通りで だいたい 2 * 107 くらいの計算量になるので間に合う。最終的にはこのDP表のうち操作回数がKに等しいか、あるいは0フラグが立っているものを足し合わせれば答えが出る。

感想

本番結構適当に辻褄を合わせて通してしまった

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, std.bitmanip;

immutable long MOD = 10^^9 + 7;

void main() {
    auto s = readln.split;
    auto N = s[0].to!int;
    auto K = s[1].to!long;
    auto A = readln.split.map!(to!long).array;
    auto cnt = new long[](N);
    long ans = 0;

    foreach (i; 0..N) {
        for (long a = A[i]; a > 0; a /= 2) {
            cnt[i]++;
        }
    }

    auto dp = new long[][][](N+1, 65*N, 2);
    //foreach (i; 0..N+1) foreach (j; 0..50*N) foreach (k; 0..2) dp[i][j][k] = -1;
    foreach (i; 0..N) if (A[i] == 0) dp[0][0][1] = 1;
    if (dp[0][0][1] == 0) dp[0][0][0] = 1;

    foreach (i; 0..N) {
        foreach (j; 0..65*N) {
            foreach (k; 0..2) {
                if (dp[i][j][k] == 0) continue;
                foreach (using; 0..cnt[i]+1) {
                    if (j + using > K) continue;
                    (dp[i+1][j+using][k||(using==cnt[i])] += dp[i][j][k]) %= MOD;
                }
            }
        }
    }

    //dp.each!writeln;
    foreach (j; 0..min(65*N, K+1)) (ans += dp[N][j][1]) %= MOD;
    if (K < 65*N) (ans += dp[N][K][0]) %= MOD;


    ans.writeln;
}