E – Yutori 解説 (AtCoder Beginner Contest 161)
問題概要
N日間のうちK日働く。
- 一日働いたら、直後のC日間働けない
- i+1 日目は、Sの i 文字目が 'o’ の時だけ働ける
必ず働く日を全て求めよ。
制約
- \(1 \leq N \leq 2\times 10^5\)
- \(1 \leq K \leq N\)
- \(0 \leq C \leq N\)
考え方
「必ず働く日」や「必ず働かなくても良い日」というのは、どのような状況で生じるのかを考えてみます。
必ず働く日が存在しない場合
例えば以下の入力の時を考えてみます。
5 1 0
ooooo
0日目から4日目までの5日間、どの日でも働くことができます。このうちの1日だけ働けばいいので、必ず働かなくてはならない日というのは生じません。
つまり、最大で5日間働けて、働く必要のあるのは 1日のみです。
これを一般化すると、最大で L 日間働けて、働く必要のあるのが K 日である時、
- L > K の時：「働かない日を1つ以上選べる」ので、必ず働かなくてはならない日はない。
と言うことができます。
Lを求める方法
働ける最大日数 L を求める方法を考えてみます。
問題設定から区間スケジューリング問題に似た雰囲気を感じますが、今回は「働けない間隔C」が固定なので、前から貪欲に働く日を決定していけば良いです。
つまり、最も早く・なるべく多く働くにはいつ働けば良いかを考えましょう。
例えば以下の入力例を考えてみましょう。
11 3 2
ooxxxoxxxoo
- 1日目は 'o’ なので働ける。その後 2 日は働けない
- 4日目は 'x’
- 5日目は 'x’
- 6日目は 'o’ なので働ける。その後 2 日は働けない
- 9日目は 'x’
- 10日目は 'o’ なので働ける。その後 2 日は働けない
となり、最大で3日間働くことができます。
どのように働く必要があるか
LとKが等しいならば、L を求める時の方法で、全体を K 個に分割することができます。それぞれの区間中では、必ず1回働く必要があります。
さきほどの例では、
ooxxx oxxx oo
と分割できます。
「必ず働かなくても良い日」というのは、区間中に他の選択肢が存在する時に生じます。
逆に「必ず働く必要がある日というのは、区間中に他の選択肢が存在しない時に生じます。
矛盾が生じないように（休息期間Cと被らないように）うまく選択する必要がありますが、実は右から働く日を貪欲に決定してやれば良いです。
以下の入力を考えてみましょう。
16 4 3
ooxxoxoxxoxoxxxo
これは以下のように分割できます。
ooxx oxoxx oxoxxx o
例えば、5日目に必ず働く必要があるかを考えてみましょう。他の選択肢として「7日目に働く」というものがありますが、仮に10日目（正確には7+1日目から7+3日目まで）に働く必要があるとすれば矛盾します。
しかし、11日目のかわりに13日目に働くことができるので今回は大丈夫です。
つまり、右側から最も遅く K 日働くならいつ働けば良いかを決定していけば、区間中に他の選択肢があるかどうかを簡単に判定することができます。
このように、条件などがあって後ろの方が選択肢が少なくなる場合は、後ろから貪欲に決定していくとうまくいくことがあります。
解法
最も早く・なるべく多く働くにはいつ働けば良いかを左から貪欲に計算する。働くことができる最大日数をLとして
- L>K のとき、必ず働く必要のある日は存在しない
そうでない時は、全体を K 個の区間に分けて、以下の2つを貪欲に計算して区間ごとに比べる。
- 左から計算 A：最も早く K 日働くにはいつ働けば良いか
- 右から計算 B：最も遅く K 日働くならいつ働けば良いか
A[ i ] == B[ i ] となった場合は、必ず働く必要があります。
一つ目は、Lを計算する時にすでに求めているので、それを使いまわせば良いです。
C++ での実装例
#include <bits/stdc++.h>
#define ALL(obj) begin(obj), end(obj)
#define debug(x) cerr << #x << ": " << x << '\n'
using namespace std;
const int INF = 2100100100;
int main() {
int N, K, C;
string S;
cin >> N >> K >> C >> S;
vector<int> A; // 左から貪欲
for (int i = 0; i < N;) {
if (S[i] == 'o') {
A.push_back(i);
i += C + 1;
} else {
i++;
}
}
int L = (int)A.size();
if (L > K) return 0;
vector<int> B(K); // 右から貪欲
A.push_back(N);
B.push_back(INF);
for (int i = K - 1; i >= 0; i--) {
for (int l = A[i + 1] - 1; l >= A[i]; l--) {
if (S[l] == 'o' && l + C < B[i + 1]) {
B[i] = l;
break;
}
}
}
vector<int> ans;
for (int i = 0; i < K; i++) {
if (A[i] == B[i]) ans.push_back(A[i]);
}
for (auto i : ans) {
cout << i + 1 << "\n";
}
}