nakashiiiの自由帳

自由に書きます

AtCoder Beginner Contest 247 E - Max Min

本番ではACできず、解説見たら区間数え上げの典型(?)としてとても勉強になったのでメモ
基本的には以下の2つの解説を(おおよそ)そのまま実装していく

勉強になったこと

  • しゃくとり法の応用
  • 区間の数え上げの時に、最大値や最小値のインデックスをうまく使う

問題

E - Max Min

前提知識

  • しゃくとり法

しゃくとり法はけんちょんさんの記事がとっても分かりやすい(私もこの記事で3週間ほど前にしゃくとりデビューしました)
しゃくとり法 (尺取り法) の解説と、それを用いる問題のまとめ - Qiita

考察

公式解説に記載のある以下のサンプルケースを考える
X=4,Y=2
A=(4,2,5,4,3,4,2)

まず、数え上げたい区間の条件は以下

<条件>

  • 区間[L,R]において、最大値はX、最小値はY となる区間を数えよ

例えば、以下のような感じ

  • 区間[0, 1] → 最大値は4、最小値は2(OK)
  • 区間[0, 2] → 最大値は5、最小値は2(NG)
  • 区間[4, 6] → 最大値は4、最小値は2(OK)

ただ、このままの条件では数え上げしにくいので、以下のように条件を言い換える 。

<条件>
以下の2つを満たす区間を数えよ

  • 区間[L,R]において、最大値はX以下、最小値はY 以上
  • 区間[L,R]において、X、Yの両方を含む

ここまでできれば、準備は完了。この問題を解けるファーストステップとして、まずはこの条件の言い換えができるか?が大切(たぶん)
あとは上記の区間をどうやって数えるか?を考える。
制約がN=105くらいあるので、LとRの全探索のO(N2)だと間に合わない。
今回はO(N)で解ける方法を2つ実装していく。

方法1:条件を満たす区間に分割して、それぞれの区間に対してしゃくとり法を繰り返す

まずはしゃくとり法を使う方法。ざっくりと数え上げ手順は以下。

  • 区間[L,R]において、最大値はX以下、最小値はY 以上」という条件を満たす列をAから抽出する。
  • 上記で抽出した列に対して、しゃくとり法で「区間[L,R]において、X、Yの両方を含む」区間を数え上げ
  • 上記2つを繰り返す

以下は上記の手順を図に示したもの。青色の区間に対してしゃくとり法を適用していく。赤色の区間(値)は条件を満たさないので無視する。

数え上げ問題の区間分割としゃくとり法の適用
区間の分割イメージ

というわけで、しゃくとり法はO(N)なので、全体でもO(N)で数え上げができて、めでたし、めでたし。という感じ。
最後に実装は以下のようになる。

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

// 変数宣言------------------
int n, x, y;

// 関数定義------------------
ll syakutori(vector<int>& a, int start, int end) {
  ll ans = 0;
  int r = start;
  int xcnt = 0;  // xの出現回数
  int ycnt = 0;  // yの出現回数

  for (int l = start; l < end; l++) {
    while (r < end && !(xcnt && ycnt)) {
      if (a[r] == x) xcnt++;
      if (a[r] == y) ycnt++;
      r++;
    }

    // x, yの両方が出現する区間だけ数え上げ
    if (xcnt && ycnt) {
      ans += (end - r + 1);
    }

    if (l == r) {
      r++;
    } else {
      if (a[l] == x) xcnt--;
      if (a[l] == y) ycnt--;
    }
  }

  return ans;
}

// メイン------------------
int main() {
  cin >> n >> x >> y;
  vector<int> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  // x <= a[i] <= y となる区間 [l,r) を考える
  ll ans = 0;
  int r = 0;
  for (int l = 0; l < n; l++) {
    // a[l] が条件満たさないのでNG
    if (a[l] > x || a[l] < y) {
      r = l + 1;
      continue;
    }

    // rを求める
    while (y <= a[r] && a[r] <= x) {
      r++;
    }

    // しゃくとりで区間を数える
    ans += syakutori(a, l, r);

    // 次に探索するlはr+1となる
    l = r;
    r++;
  }

  cout << ans << endl;
}

方法2:Xのインデックス、Yのインデックス、NG数値のインデックスを持ちながら、0〜N-1まで順次数え上げ

個人的にはこちらの方法の理解が難しかった。でも実装はこっちの方がグッとシンプルになるし、自分の手札としてもっておけると強そうだと思った。
基本方針としては、以下のようになる

  • 区間の右端Rを0〜N-1まで走査する
  • その際、以下の3つを順次更新していく
    • 区間内で最も右に登場するX
    • 区間内で最も右に登場するY
    • 区間内で最も右に登場するNG値(= Xより大きい or Y未満)
  • 上記のそれぞれのインデックスをRx, Ry, Rngとすると、(Rng < Rx and Rng < Ry)となる時、条件に合致する区間なので数え上げる。add = min(Rx, Ry) - Rng が条件を満たす区間[L, R]での区間の個数

上記を図示したのが以下。R=2の場合とR=6の場合の2つ。
(図ではRが小文字rになってます、ごめんなさい)

まずは R=2 の場合(Rng < Rx and Rng < Ry の条件NGで、数え上げ区間なし)

最大値と最小値のインデックスを利用した数え上げ問題のイメージ
R=2 の場合

次に R=6 の場合(Rng < Rx and Rng < Ry の条件OKで、数え上げ区間あり)

最大値と最小値のインデックスを利用した数え上げ問題のイメージ
R=3 の場合

よって、Rの走査はO(N)、区間のカウントはO(1)でできるので、全体ではO(N)になってめでたし、めでたしとなる。
実装は以下。

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

// 変数宣言------------------
int n, x, y;

// 関数定義------------------

// メイン------------------
int main() {
  // デバッグ用
  ifstream in("input.txt");
  if (in.is_open()) {
    cin.rdbuf(in.rdbuf());
  }

  cin >> n >> x >> y;
  int a[n];
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  int rx = -1;  // xが出てくる一番右のインデックス
  int ry = -1;  // yが出てくる一番右のインデックス
  int rng = -1;  // xより上 or y未満 が出てくる一番右のインデックス

  ll ans = 0;
  for (int r = 0; r < n; r++) {
    if (a[r] == x) rx = r;
    if (a[r] == y) ry = r;
    if (a[r] < y || x < a[r]) rng = r;
    if (rng < rx && rng < ry) ans += (min(rx, ry) - rng);
  }

  cout << ans << endl;
}

感想

シンプルにしゃくとり法を適用するだけの問題は解いたことがあったが、今回のようにしゃくとり法が適用できるように前処理してから、しゃくとりするのは初めてだったので、勉強になった。
あと、最大値や最小値のインデックスをうまく使うとシンプルなコードで数え上げできたりするのが、とてもビックリした。何回か使いながら、自然と考察で思いつけるようになりたいな〜。