【AtCoder】ABC 419 E - Subarray Sum Divisibility | 緑コーダーが解くAtCoder

atcoder.jp

配点: 475 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 1298 / NoviSteps: 1D

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

問題概要

長さ $N$ の整数列 $A$ が与えられる。これから、以下の操作を繰り返し行うことにより、 $A$ のすべての長さ $L$ の連続部分列についてその総和が $M$ の倍数であるようにしたい。

$1 \leq i \leq N$ なる整数 $i$ を選び、 $A_i$ の値を $1$ 増やす。

このとき、操作回数として考えられる最小値を求めよ。

制約

$1 \leq N, M \leq 500$
$1 \leq L \leq N$
$0 \leq A_i \lt M$
入力される値はすべて整数

考察

条件の言い換え

まず、「 $A$ のすべての長さ $L$ の連続部分列についてその総和が $M$ の倍数である」という条件を言い換えよう。

連続部分列の初項を $A_l$ としたとき、とりあえず $\displaystyle \sum_{i=l}^{l+L-1} A_i \equiv 0 \pmod M$ である必要がある。

また、この初項を1つ右にずらしたとしても、 $\displaystyle \sum_{i=l+1}^{l+L} A_i \equiv 0 \pmod M$ が成り立たなければならない。

ここで、1つ目の式から2つ目の式を引くと、

$\begin{align*} A_i - A_{i+L} \equiv 0 \iff A_i \equiv A_{i+L} \pmod M \quad (1 \leq i \leq N-L) \end{align*}$

が得られる。

これは、「インデックスが $L$ 離れた $A$ の要素同士は、 $M$ による剰余が等しくなければならない」と解釈できる。

したがって、 $A$ の初めの $L$ 個の要素それぞれについて、 $M$ による剰余をどの値にするか(操作回数)を決めたら、残りの $N - L$ 個の要素に対する操作回数も自動的に確定することになる。

A の各要素をグループ分けして、コストテーブルを前計算する

上で言い換えた条件を基に、 $A$ の各要素を以下のようにグループ分けすることを考える。

グループ $1$ : $A_1, A_{L+1}, A_{2L+1}, \dots,$
グループ $2$ : $A_2, A_{L+2}, A_{2L+2}, \dots,$ ...
グループ $L$ : $A_L, A_{2L}, A_{3L}, \dots,$

グループ $i$ の要素について、 $M$ による剰余を $m_i \: (0 \leq m_i \leq M-1)$ に統一することにする。

ただし、 $\displaystyle \sum_{i=1}^{L} A_i \equiv \sum_{i=1}^{L} m_i \equiv 0 \pmod M$ である。

このとき、

$\mathrm{cost}_{i, j} :=$ $A_i, A_{i+L}, A_{i+2L}, \dots,$ を $M$ による剰余を全て $j$ にするために必要な操作回数(コスト)の最小値

と定義すれば、これは

- ただし、 $S_i = \{k \mid 1 \leq k \leq N, k \equiv i \pmod L\}$

と計算できる。操作では、要素に加算することしかできないので、 $j - A_k \lt 0$ のときは追加で $M$ を加算する必要があることに注意しよう。

動的計画法で答えを求める

さて、ここまでの議論を踏まえると、我々がやりたいのは

個にグループ分けしたグループ内で、統一するによる剰余の値をそれぞれ決めていく。ただし、以下の条件を満たすように決める。
- 決めた剰余の値の合計が $M$ の倍数になる。
- 操作回数が最小になる。

ということである。グループ毎に統一する $M$ による剰余の値を逐次的に決めていけるので、これは動的計画法の出番だ。

DPテーブルは以下のように定義できる。

$\mathrm{dp}_{i, j} :=$ $i$ 個目のグループまでについて、統一する $M$ による剰余の値を決めたとき、その剰余の和を $M$ で割った余りが $j$ となるような組み合わせの中での、コストの合計の最小値

初期値は $\mathrm{dp}_{0, 0} = 0$ であり、他は全て $\infty$ とできる。

グループ $i$ までが確定していて、新たにグループ $i+1$ で $M$ による剰余の値を $k$ に統一すると決めたなら、コストは $\mathrm{cost}_{i, k}$ 必要であり、グループ $i+1$ までの「統一する剰余の値の和を $M$ で割った余り」は、 $(j + k) \mod M$ である。

したがって、遷移式は以下のようになる。

$\mathrm{dp}_{i+1, (j+k) \mod M} \leftarrow \min(\mathrm{dp}_{i+1, (j+k) \mod M}, \mathrm{dp}_{i, j} + \mathrm{cost}_{i, k})$

求める答えは $\mathrm{dp}_{L, 0}$ を参照すればよい。

計算量について。

コストテーブルは $i = 1, 2, \dots, L, \: j = 0, 1, \dots, M-1$ の二重ループ内で、要素数約 $\frac{N}{L}$ 個に対してのループを行っているので、 $O(NM)$ と評価できる。

また、DPテーブルは $O(LM^2)$ で計算が可能である。

以上から、この問題を $O(LM^2)$ で解くことができた。

実装例

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
constexpr int INF = 1e+9;

template <typename T>
inline bool chmin(T &a, T b) { return ((a > b) ? (a = b, true) : (false)); }

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

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

    vector<vector<int>> cost(L, vector<int>(M, 0));
    rep(i, 0, L) rep(j, 0, M)
    {
        int sum = 0;
        for (int k = i; k < N; k += L)
        {
            sum += (j - A[k] + M) % M;
        }
        cost[i][j] = sum;
    }

    vector<vector<int>> dp(L + 1, vector<int>(M, INF));
    dp[0][0] = 0;
    rep(i, 0, L) rep(j, 0, M) rep(k, 0, M)
    {
        chmin(dp[i + 1][(j + k) % M], dp[i][j] + cost[i][k]);
    }

    cout << dp[L][0] << endl;

    return 0;
}