[CSA] Round #2 Circular Subarrays

問題

長さが \( N \) のリング状の配列 \( a \) の長さが \( K \) の部分配列の和がすべて等しくなるような操作回数の最小値を求めます。

方針

部分配列の和

円形の配列 \( a \) が与えられたとき、サイズ \( K \) の連続する部分配列は \( N \) 個あり、その部分配列の和が全て等しくなるように配列の要素を増減させます。どういうことかというと、 \( i \) 番目の部分配列の和を \( S_i \) とすると、
\[
\begin{eqnarray}
S_1 &=& a_1 + a_2 + \cdots + a_K \\
S_2 &=& a_2 + a_3 + \cdots + a_{K + 1}\\
& \vdots &\\
S_{n-1} &=& a_{n – 1} + a_{n} + \cdots + a_{K-2}\\
S_n &=& a_n + a_1 + \cdots + a_{K – 1}
\end{eqnarray}
\]
これらの和が全て等しくなるように操作を行います。全ての和が等しいような数列 \( a \) が満たすべき条件を考えると、\( S_1 = S_2 \) より、\( a_1 = a_{K + 1}\) であることが分かります。同様にしてすべての部分配列に対して条件を考えると、値が等しくなるような要素のグループが出来上がります。
公式の解説では、このようなグループの数は、
\[ \dfrac{N}{gcd(N, K)} \]
となるようです。また、このことは解答を作成するうえでは重要ではありません。したがって、グループの数を \( m \) とし、その \( i \) 番目の集合を \( G_i \) とすると、
\[
\begin{eqnarray}
G_1 &=& \{a_1, a_{1 + K}, a_{1 + 2K}, \cdots \}\\
G_2 &=& \{a_2, a_{2 + K}, a_{2 + 2K}, \cdots \}\\
& \vdots &\\
G_m &=& \{a_m, a_{m + K}, a_{m + 2K}, \cdots \}\\
\end{eqnarray}
\]
のようになります。ここで注意すべき点は、あるグループに属する配列の要素番号は他のグループには属しません。
各グループの要素の値を等しくするために必要な操作の最小回数は、そのグループの中央値と要素の絶対値の和となります。この考えは割とよく見かけると思います。
したがって、全てのグループに対して必要な操作回数の和が答えとなります。

コード

提出したコード

グループ分けとコストの計算

  int ans = 0;
  set<int> s;
  for (int i = 0; i < N; i++) {
    vector<int> v;
    int idx = i;
    if (s.count(idx) != 0) continue;
    while (s.count(idx) == 0) {
      v.push_back(a[idx]);
      s.insert(idx);
      idx += K;
      idx %= N;
    }
    sort(v.begin(), v.end());
    int mid = v[v.size() / 2];  
    for (int j : v) {
      ans += abs(j - mid);
    }
  }