【AtCoder】ABC 406 C - ~ | 緑コーダーが解くAtCoder

atcoder.jp

配点: 350 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: ??? / NoviSteps: 1Q

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

問題概要

整数列 $A$ に対し、 $A$ が「チルダ型」とは以下の4つの条件を全て満たすことであると定義する。

$|A| \geq 4$
$A_1 \lt A_2$
$A_{i - 1} \lt A_i \gt A_{i + 1}$ を満たす $2 \leq i \lt |A|$ なる整数 $i$ はちょうど1個である
$A_{i - 1} \gt A_i \lt A_{i + 1}$ を満たす $2 \leq i \lt |A|$ なる整数 $i$ はちょうど1個である

数列 $(1, 2, \dots, N)$ を並べ替えて得られる数列 $P = (P_1, P_2, \dots, P_N)$ が与えられるので、 $P$ の連続部分列であってチルダ型である数列の個数を求めよ。

制約

$4 \leq N \leq 3 \times 10^5$
入力される値は全て整数

考察

とりあえず、数列 $P$ の下に不等号を書いて、チルダ型となる数列の範囲と共に可視化してみる(以下では、 $P_i$ の要素はこの不等号とみなす)。

これを参考にすると、数列 $P$ のある連続部分列 $(P_l, P_{l+1}, \dots, P_r)$ がチルダ型となるためには、次の条件が満たされていることだと分かる。

$P_l$ は<である
連続部分列内に<>と><がただ1か所存在する

したがって、以下のような探索アルゴリズムを考えることができる。

$P_1$ から $P_N$ まで順に不等号を見ていき、 $P_i =$ <となるところを見つけ、 $l = i$ とする。
$j = l, l+1, \dots$ として、 [tex: P_j = >] となる場所と、再び $P_j =$ <となる場所を見つける。
1. の後者が見つかったら、答えのカウンタを1増やして、 $i = l + 1$ から再び 1. を行う。

しかし、これは計算量が $O(N^2)$ であるため、本問の制約下では TLE してしまう。どうすればよいだろうか。

先ほどの図をよく見てみると、チルダ型となっている部分の不等号の並び順は、以下の要素を順に連結したものであることが分かる。

1個以上の<
1個以上の>
1個以上の<

ここで、前半の<の個数を $x$ 個、後半の>の個数を $y$ とすると、チルダ型の始点 $P_l$ を固定したときに終点 $P_r$ は $y$ 通り存在し、 $P_l$ の動かし方は $x$ 通り存在することから、この区間には全部で $xy$ 個のチルダ型が存在することが分かる。これはサンプルケース3がいい例となっている。

したがって、「1個以上の<」、「1個以上の>」、「1個以上の<」が連続している区間を探してそれぞれの個数を求めれば、この区間に含まれるチルダ型は $O(1)$ で計算できるため、全体の探索は $O(N)$ の計算量まで落とすことができた。

<や>が連なる個数を求めるためには、ランレングス圧縮が有効である。

algo-logic.info

詳しくは実装例を参照のこと。ただし、実装例では不等号を書く代わりに、新たに長さ $N - 1$ の数列 $B$ を以下のように定義して、それに対してランレングス圧縮を行っている。

$B_i = \left\{ \begin{array}{ll} 0 & (P_i \lt P_{i+1}) \\ 1 & (P_i \gt P_{i+1}) \end{array} \right.$

※ <を $0$ に、>を $1$ に置き換えているということ。

実装例

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

#define all(x) (x).begin(), (x).end()
#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

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

std::vector<std::pair<char, int>> encode(const std::string &str)
{
    int n = (int)str.size();
    std::vector<std::pair<char, int>> res;
    for (int l = 0; l < n;)
    {
        int r = l + 1;
        for (; r < n && str[l] == str[r]; r++)
        {
        };
        res.push_back({str[l], r - l});
        l = r;
    }
    return res;
}

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

    vector<int> B(N - 1);
    rep(i, 0, N - 1)
    {
        if (P[i] < P[i + 1])
            B[i] = 0;
        else
            B[i] = 1;
    }

    vector<pair<char, int>> rle = encode(string(all(B)));
    int l = rle.size();

    ll ans = 0;
    rep(i, 0, l - 2)
    {
        if (rle[i].first == 0 && rle[i + 1].first == 1 && rle[i + 2].first == 0)
        {
            ans += (ll)rle[i].second * rle[i + 2].second;
        }
    }

    cout << ans << endl;

    return 0;
}