AtCoder Beginner Contest D - Range Add Query
問題はこちら
回答
効率的な解き方が分からず、解説を見た
以下の不変量に着目して解いていた
実際は、クエリでl,rが与えられるため、mod K毎にインデックス毎に累積和をとっておく必要
#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>
#include <iomanip>
#include <random>
int N, K, Q;
std::vector<long long> A;
int main()
{
std::cin.tie(0);
std::ios::sync_with_stdio(false);
std::cin >> N >> K;
A.resize(N + 1);
for (int i = 1; i <= N; i++) {
std::cin >> A[i];
}
std::cin >> Q;
// modKでグループ分けしたときのAの値の累積和を求める
// x方向は、1-indexed
std::vector<std::vector<long long>> modK = std::vector< std::vector<long long>>(K, std::vector<long long>(N + 1));
for (int i = 1; i <= N; i++) {
// iに対して累積和
int m = i % K;
for (int k = 0; k < K; k++) {
if (k != m) {
// 引き継ぐ
modK[k][i] = modK[k][i - 1];
}
}
//累積和
modK[m][i] = modK[m][i - 1] + A[i];
}
std::vector<std::string> ans;
while (Q--) {
int l, r;
std::cin >> l >> r;
//区間l~rの各modKグループの総和を求める
long long modKLR;
bool res = true;
for (int k = 0; k < K; k++) {
// modK[k]のl~rまでの和
modKLR = modK[k][r] - modK[k][l - 1];
// 各グループの総和が等しいことを確認する
long long val = modK[0][r] - modK[0][l - 1];
if (val != modKLR) {
res = false;
break;
}
}
if (res) {
ans.push_back("Yes");
}
else {
ans.push_back("No");
}
}
for (int i = 0; i < ans.size(); i++) {
std::cout << ans[i] << std::endl;
}
return 0;
}
補足
以下の不変量に着目して解くこともできる
コードは以下で、先ほどの通り、クエリでl,rが与えられるため、mod K毎にインデックス毎にD[i]の累積和をとっておく必要がある
ただ、全体の累積和から区間l,rの累積和を取得するための境界処理が面倒
#include <iostream>
#include <vector>
using namespace std;
int main() {
// 入出力の高速化
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, K;
if (!(cin >> N >> K)) return 0;
vector<long long> A(N + 1, 0);
for (int i = 1; i <= N; i++) {
cin >> A[i];
}
// 差分配列 D を構築 (D[i] = A[i] - A[i-1])
vector<long long> D(N + 2, 0);
for (int i = 1; i <= N; i++) {
D[i] = A[i] - A[i - 1];
}
// prefD[m][i] : 1番目からi番目までの要素のうち、インデックス % K == m である D の和
vector<vector<long long>> prefD(K, vector<long long>(N + 1, 0));
for (int i = 1; i <= N; i++) {
for (int m = 0; m < K; m++) {
prefD[m][i] = prefD[m][i - 1];
}
prefD[i % K][i] += D[i];
}
int Q;
cin >> Q;
while (Q--) {
int l, r;
cin >> l >> r;
bool is_good = true;
// 部分列 A[l...r] から作られる差分配列は以下の要素を持つ
// 先頭: A[l] （元のインデックス l 相当）
// 中間: D[i] （l < i <= r までの各 D[i]）
// 末尾: -A[r] （元のインデックス r+1 相当）
for (int m = 0; m < K; m++) {
// 中間部分の和
long long current_sum = prefD[m][r] - prefD[m][l];
// 境界（先頭・末尾）がこのグループ(余りm)に属していれば足し引きする
// 境界（先頭・末尾）がこのグループ(余りm)に属しているときにちょうどmodKグループの累積和が変化するため
// (★)
// prefD[m][r] - prefD[m][l]では、D[l+1] ~ D[r]の和となる
// そのため、D[l] ~ D[r]の和を求めるためには、prefD[m][r] - prefD[m][l]にD[l]を足す必要
// D[l]=A[l]-A[l-1]となる
// しかし、範囲外であるA[l-1]は0なため、A[l]のみを加算すればよい
if (l % K == m) {
current_sum += A[l];
}
//
// D[i], D[i＋K]において、i~i+(K-1)の長さKの区間に加算された際、
// D[i]にc, D[i＋K]に-cが加算されることを考慮
// ※cが加算されたとしても、D[i]とD[i＋K]の和は依然として0
// ※この性質により、インデックスに対するmodKによって、D[i]をグルーピングした際、
// 同一グループ内の値の総和は0
// 今回、同じmodKでD[i]をまとめていることを考えると、i+Kなのか、i+xKなのかは重要ではない
// そのため、今注目しているmodK(ループ変数m)と同じmodKのインデックスi+xKは、
// r以下の中で最も大きいインデックスとする
// その時、インデックスi+xKがちょうどr+1の時、r+1まで巻き込んだD[i]の総和を
// 求めなければ、l~rの正確な和は得られない
// そのため、prefD[m][r] - prefD[m][l]にD[r+1]を足せばよい
// D[r+1] = A[r+1] - A[r]だが、範囲外であるA[r+1]は0
// よって、prefD[m][r] - prefD[m][l]に-A[r]を足せばいい
//
if ((r + 1) % K == m) {
current_sum -= A[r];
}
// 階差配列の考え方では、すべての余りグループの総和が 0 である必要がある
if (current_sum != 0) {
is_good = false;
break;
}
}
if (is_good) cout << "Yes\n";
else cout << "No\n";
}
return 0;
}
上記コードの★の部分は必要に応じて以下の図を用いて理解すること
