AtCoder が提供しているABC（AtCoder Beginner Contest）337 のD問題をC++とPythonで解いてみました。ABC337は、2024年1月20日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Cheating Gomoku Narabe（Difficulty : 760）
問題はリンク先をご覧ください。
ABC337 D問題 Cheating Gomoku Narabe
行または列毎に条件を満たすか調べます。それぞれの行では累積和を求めて判定します。AtCoder Problems による Difficulty は 760 でした。
解答案
それぞれの行（または列）では、K 個の連続する要素について、以下の手順で調べます。
- K 個の要素のひとつでも ‘x’ があれば、対象外となる。
- K 個の要素が すべて’.’ または ‘o’ の場合、’.’ の個数をポイントとする。
- ポイントの最小値を更新する。
2種類の累積和を計算しておけば、K 個の要素に ‘x’ が含まれているか、Ｋ 個の要素の中の ‘o’ の個数を $O(1)$ で求めることができます。
C++ プログラム例（ABC337D）
プログラムを補足します。
- 行の文字列を s、列の文字列を t とします（38ー47行目）。
- 関数 check で、それぞれの行または列のポイントの最小値を求めます。
- 文字 ‘x’ の個数の累積和 sum_x、文字 ‘.’ の個数の累積和 sum_none を求めておきます。累積和の処理は、定番コードとなっています。
- 連続する k 個に ‘x’ を含まない場合（26行目）、文字 ‘.’ の個数をポイントとして、そのポイントの最小値を返します（27行目）。
- 変数 result が初期値の場合、-1 にします。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
#define INF (1 << 30)
int check(string s, int n, int k)
{
int result = INF;
vector<int> sum_x(n + 1);
vector<int> sum_none(n + 1);
sum_x[0] = 0;
sum_none[0] = 0;
for (int i = 0; i < n; ++i) {
sum_x[i + 1] = sum_x[i];
if (s[i] == 'x') {
++sum_x[i + 1];
}
sum_none[i + 1] = sum_none[i];
if (s[i] == '.') {
++sum_none[i + 1];
}
}
for (int i = 0; i + k <= n; ++i) {
if (sum_x[i + k] - sum_x[i] == 0) {
result = min(result, sum_none[i + k] - sum_none[i]);
}
}
return result;
}
int main()
{
int h, w, k;
cin >> h >> w >> k;
vector<string> s(h);
for (int i = 0; i < h; ++i) {
cin >> s[i];
}
vector<string> t(w);
for (int i = 0; i < w; ++i) {
for (int j = 0; j < h; ++j) {
t[i] += s[j][i];
}
}
int result = INF;
for (int i = 0; i < h; ++i) {
result = min(result, check(s[i], w, k));
}
for (int i = 0; i < w; ++i) {
result = min(result, check(t[i], h, k));
}
if (result == INF) {
result = -1;
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC337D）
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 337 D"""
INF = 1 << 30
def check(s, n, k):
result = INF
sum_x = [0] * (n + 1)
sum_none = [0] * (n + 1)
for i in range(n):
sum_x[i + 1] = sum_x[i]
if s[i] == 'x':
sum_x[i + 1] += 1
sum_none[i + 1] = sum_none[i]
if s[i] == '.':
sum_none[i + 1] += 1
for i in range(n - k + 1):
if sum_x[i + k] - sum_x[i] == 0:
result = min(result, sum_none[i + k] - sum_none[i])
return result
h, w, k = map(int, input().split())
s = [input() for i in range(h)]
t = ["" for i in range(w)]
for i in range(w):
for j in range(h):
t[i] += s[j][i]
result = INF
for i in range(h):
result = min(result, check(s[i], w, k))
for i in range(w):
result = min(result, check(t[i], h, k))
if result == INF:
result = -1
print(result)こちらも「AC」と判定されました。
最後に
それぞれの行（または列）では、累積和で判定すると読みやすいプログラムになりました。
問題を読み、最初に書いたプログラムは以下です。AC 判定となります。記事で紹介したコードとの差分は関数 check の実装だけです。しゃくとり法を意識していますが、ロジックが冗長になっていると感じて、公式解説に従い、累積和を使ったプログラムを紹介しました。
#include <bits/stdc++.h>
using namespace std;
#define INF (1 << 30)
int check(string s, int n, int k)
{
int result = INF;
int left = 0;
int right = 0;
int point = 0;
while ((left < n)&&(right < n)) {
right = left;
point = 0;
bool check_next = false;
for (int i = 0; i < k; ++i) {
if (right >= n) {
check_next = true;
break;
}
if (s[right] == 'x') {
left = right + 1;
check_next = true;
break;
}
if (s[right] == '.') {
++point;
}
++right;
}
if (check_next) {
continue;
}
result = min(result, point);
while (1) {
if (right >= n) {
break;
}
if (s[right] == 'x') {
left = right + 1;
break;
}
if (s[right] == '.') {
++point;
}
if (s[left] == '.') {
--point;
}
result = min(result, point);
++left;
++right;
}
}
return result;
}
int main()
{
int h, w, k;
cin >> h >> w >> k;
vector<string> s(h);
for (int i = 0; i < h; ++i) {
cin >> s[i];
}
vector<string> t(w);
for (int i = 0; i < w; ++i) {
for (int j = 0; j < h; ++j) {
t[i] += s[j][i];
}
}
int result = INF;
for (int i = 0; i < h; ++i) {
result = min(result, check(s[i], w, k));
}
for (int i = 0; i < w; ++i) {
result = min(result, check(t[i], h, k));
}
if (result == INF) {
result = -1;
}
cout << result << endl;
return 0;
}
