【AtCoder】ABC 390 D - Stone XOR | 緑コーダーが解くAtCoder

atcoder.jp

配点: 400 点 / 実行時間制限: 3 sec / メモリ制限: 1024 MB / Difficulty: 1607

問題概要
制約
考察
実装例
コメント

問題概要

$1$ から $N$ までの番号が付けられた $N$ 個の袋がある。はじめ、袋 $i$ には $A_i$ 個の石が入っている。

これから「2つの袋 $A, B$ を選び、袋 $A$ に入っている石をすべて袋 $B$ に入れる」という操作を $0$ 回以上繰り返すことを考える。

操作を繰り返した後の状態における次の値としてあり得るものが何個あるか求めよ。

袋 $i$ に入っている石の個数を $B_i$ としたときの、 $B_1\oplus B_2\oplus\cdots\oplus B_N$ の値。ただし、 $\oplus$ は排他的論理和を表す。

制約

$2\leq N \leq 12$
$1\leq A_i \leq 10^{17}$
入力はすべて整数。

考察

操作を行うごとに石は1つの袋に集約されていき、最終的にはいくつかの袋にまとめられる。

したがって、同じグループ内の袋は1つにまとめるとして、 $N$ 個の袋をいくつかのグループに分けることを考えよう。

本問は $N$ の制約が非常に小さく、グループ分けの全探索が可能そうだ。

一般に、 $n$ 個のものを区別のない $k$ 個以下のグループに分ける場合の数を「ベル数」といい、 $B(n)$ で表す。 $B(0) = B(1) = 1, \: B(2) = 2, \: B(3) = 5, \: \cdots, \: B(12) = 4213597$

※ ベル数自体を覚えておく必要はないが、 $N$ の特殊な制約から、全探索が可能なパターン数になるんだろうなという感覚は持っておきたいかも。

さて、グループ分けの全探索をどのように行うかだが、ここで下図のような樹形図を考える。なお、丸囲みの数字は袋を表し、点線で囲まれた袋は同じグループに含まれることを示す。

階層を下るごとに新たな袋の追加が行われ、左から順に既存のグループに「入れる」場合と「入れない」場合とで以降の状態が分岐している。最も右の状態は、追加された袋単体で新たなグループを作っている。

これをアルゴリズムに起こすと、以下のような再帰的な処理が書ける。

$\mathrm{groups}_i :=$ $i$ 番目のグループに入っている石の個数
$\mathrm{idx} :=$ 現在考えている袋の番号
$v :=$ $\mathrm{groups}$ 内の要素の XOR 和を管理しておく連想コンテナ

とする。

$\mathrm{groups} = \varnothing, \: \mathrm{idx} = 1$ で初期化。
$\mathrm{idx} = N$ ならば、 $\mathrm{groups}$ 内の要素の XOR 和を取って $v$ に追加し、処理を打ち切る( $\mathrm{idx} \gt 1$ なら $\mathrm{idx}$ を $1$ 減算して 3-1-3. または 3-4. へ戻る)。
そうでないとき、の要素数をとしてならば、
1. について、
  1. $\mathrm{groups}_{i}$ に $A_i$ を加算する(既存のグループに追加)。
  2. $\mathrm{idx} \leftarrow \mathrm{idx} + 1$ として、 2. からやり直す。
  3. $\mathrm{groups}_{i}$ に $A_i$ を減算する。
2. $\mathrm{groups}$ の末尾に $(A_i)$ を追加する(新たなグループを作成)。
3. $\mathrm{idx} \leftarrow \mathrm{idx} + 1$ として、 2. からやり直す。
4. $\mathrm{groups}$ の末尾を削除する。

上の樹形図を DFS で探索している感じ。

実装時は、 $v$ のデータ構造として C++ ではunordered_setを使うように注意すること。setだと TLE してしまう。

また、今回は DFS を再帰ラムダ関数を用いて実装した。少々書き方に癖はあるが、関数外の変数を参照できるようになるので実装がラクになる。

penguinshunya.hatenablog.com

実装例

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

using ll = long long;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

int main()
{
    int N;
    cin >> N;
    vector<ll> A(N);
    rep(i, 0, N) cin >> A[i];

    auto dfs = [&](auto dfs, vector<ll> &groups, unordered_set<ll> &st, int idx) -> void
    {
        if (idx == N)
        {
            ll sum = 0;
            for (auto &&x : groups)
            {
                sum ^= x;
            }

            st.insert(sum);
            return;
        }

        rep(i, 0, groups.size())
        {
            groups[i] += A[idx];
            dfs(dfs, groups, st, idx + 1);
            groups[i] -= A[idx];
        }

        groups.push_back(A[idx]);
        dfs(dfs, groups, st, idx + 1);
        groups.pop_back();

        return;
    };

    unordered_set<ll> st;
    vector<ll> B;

    dfs(dfs, B, st, 0);

    cout << st.size() << endl;

    return 0;
}