CODE THANKS FESTIVAL 2017 F Limited Xor Subset

問題

atcoder.jp

問題概要

$N$ 個の正の整数が与えられる。
そのうちから0個以上の整数を選んで排他的論理和を取った値が $K$ になる選び方の総数を求めよ。

考察 and 解説

DPで解く。

比較的に思いつきやすい遷移としては

$dp(i,j) := i$ 番目の整数まで見たときに排他的論理和が $j$ になる選び方の総数

があるが、 $O(Nmax(A))$ となりTLEしてしまう。

// C++
// O(Nmax(A))

for (int i = 0; i < n; ++i) {
  for (int j = 0; j <= max(A); ++j) {
    dp[i+1][j] += dp[i][j]; // xorをしない
    dp[i+1][j ^ a[i]] += dp[i][j]; // xorをする
  }
}

そのため

$a_{1} + a_{2} + ... + a_{n} \le 10^{5}$

という制約をうまく使ってみる。

ある $2$ つの整数 $t, s$ の排他的論理和は $max(t,s) * 2$ 以下になる。(大きい方の整数の最上位bitより大きいbitは変化しないから)

これを利用して、数列 $A$ を昇順にソートしてから最初に挙げた $dp$ を使ってみる。

すると遷移時には $a_{i} * 2$ までしか見なくてよくなるので、計算量は $O(N \sum a_{i})$ まで落ちる。

同時に空間計算量も工夫する必要がある。最初に $N * max(A) * 2$ の領域の配列を用意するとMLEする。

$dp(i+1,j)$ の更新に必要な情報は $dp(i,j)$ だけなので、 $2$ つの配列を用意し計算時にswapするだけで空間計算量も $O(N)$ まで落とせる。

f:id:ryo_seriru:20190506223306p:plain

提出コード

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
#include <functional>
#include <queue>
#include <string>
#include <cstring>
#include <numeric>
#include <cstdlib>
#include <cmath>
using namespace std;

typedef long long ll;

#define INF 10e17 // 4倍しても(4回足しても)long longを溢れない
#define rep(i,n) for(int i=0; i<n; i++)
#define rep_r(i,n,m) for(int i=m; i<n; i++)
#define END cout << endl
#define MOD 1000000007
#define pb push_back
#define sorti(x) sort(x.begin(), x.end())
#define sortd(x) sort(x.begin(), x.end(), std::greater<int>())
#define debug(x) std::cerr << (x) << std::endl;
#define roll(x) for (auto itr : x) { debug(itr); }

template <class T> inline void chmax(T &ans, T t) { if (t > ans) ans = t;}
template <class T> inline void chmin(T &ans, T t) { if (t < ans) ans = t;}

int main() {
  int n, k;
  cin >> n >> k;
  vector<ll> a(n);
  rep(i, n) cin >> a[i];
  vector<ll> dp_pre(200010, 0), dp_aft(200010, 0); // 2つ用意する！
  dp_pre[0] = 1; // (0,0)は1に初期化
  sorti(a); // 昇順にソート
  int limit = 0;
  for (int i = 0; i < n; ++i) {
    for (int j = 0; j <= limit; ++j) {
      dp_aft[j] = 0;
    }
    for (int j = 0; j <= limit; ++j) {
      dp_aft[j] += dp_pre[j] % MOD;
      dp_aft[j ^ a[i]] += dp_pre[j] % MOD;
    }

    limit |= a[i];
    swap(dp_pre, dp_aft); // ここで代入演算子を使うと一時オブジェクトが作られるためTLEする。std::swapはmoveを使うので高速。
  }

  cout << dp_pre[k] % MOD << endl;
}