AtCoder Beginner Contest 134 D　Preparing Boxes

表題通りです。

D Preparing Boxes

問題

atcoder.jp

問題概要

$N$ 個の数字が書かれた箱と数列 $a$ が与えられる。 $i$ の倍数の箱に入っているボールの数の合計の $mod 2$ が $a_{i}$ と一致するようにしたい。ボールの入れ方を答えよ。

考察

後ろから条件を満たすようにボールを入れていく。

例えば、 $N = 6, a = {1,0,0,0,1,1}$ とすると、 $i = 4~6$ の場合はその倍数が $N \leq 6$ に存在しないので $a_{i}$ の通りボールを入れれば良い。

次に $i = 3$ の場合、 $3$ の倍数は $6$ しかない。 $6$ には $a_{6} = 1$ よりボールが入っているので、箱3にはボールを入れる必要はない。

このように $i = 2$ のときも $N \leq 6$ の範囲で $2$ の倍数の箱のボールの数を数えていき帳尻を合わせればよい。

計算量が非常に重要で一見 $O(n^{2})$ に感じるが、実際は $O(nlogn)$ になり十分高速である。（調和級数で調べると良いらしい）

この問題から得られる知見 前から箱のボールを決めようとすると後ろの状況が未知なので混乱するが、後ろから答えを決めることで1つが決まれば前の箱の状況は一意に求まる。

コード

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

#define rep(i,n) for(long long i=0; i<n; i++)

int main() {
  int n;
  cin >> n;
  vector<int> a(n), b(n, 0);
  
  set<int> st;
  rep(i, n) cin >> a[i];

  for (int i = n; i > 0; --i) {
    int cnt = 0;
    for (int j = i + i; j <= n; j += i) {
      cnt += b[j - 1];
    }
    if (cnt % 2 != a[i-1]) {
      b[i-1] = 1;
      st.insert(i);
    }
  }

  cout << st.size() << endl;
    if (st.size()) {
      for (auto&& itr : st) {
      cout << itr << " ";
    }
  }
  

}