【AtCoder】ABC 352 D - Permutation Subsequence | 茶コーダーが解くAtCoder

実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 714

問題概要

$1, 2, \dots, N$ の順列 $P$ が与えられる。長さ $K$ の正整数列 $(i_1, i_2, \dots, i_K)$ であって、以下の条件を満たすものを「良い添字列」と定義する。すべての添字列における $i_K - i_1$ の最小値を求めよ。

定義 :

$1 \leq i_1 < i_2 < \dots < i_K leq N$ を満たす。
数列 $(P_{i_1}, P_{i_2}), \dots, P_{i_K}$ はある連続する $K$ 個の整数の順列である。

制約

入力はすべて整数。
$1 \leq K \leq N \leq 2 \times 10^5$

考察

サンプルケース3を例に考察を進めてみよう。

まず、 $(P_{i_1}, P_{i_2}, \dots, P_{i_K})$ は連続する $K$ 個の整数列になる必要があるので、考えやすいように配列 $v_j$ に「 $P_i = j$ となるような $i$ 」を格納しよう。すると、 $P, v$ の中身はそれぞれ下図のようになる。

ここで、問題文中の $a$ を $1$ から $N - K + 1$ まで増やしたとき、 $i_K - i_1$ の値がどう変わるかを見ていくと、下図のようになる。

$a$ を進めるごとに $v$ の中での探索範囲が右に一つずつずれていっていることが分かる。そして、 $i_K - i_1$ の値は探索範囲内の要素の集合を $v'_a$ として、 $\max v'_a - \min v'_a$ となっている。

つまり、求めるべき値は $\max v'_a - \min v'_a (a = 1, 2, \dots, N-K+1)$ の最小値である。

これを高速に求めるにはどうしたらよいであろうか。上の二つの図を見比べると、 $a$ を一つ進めるごとに $v'_a$ の中から要素 $v_a$ が追い出され、新たに $v_{a+K}$ が追加されて $v'_{a+1}$ を構成していることに気付くはずだ。したがって、「配列の要素の削除・追加」と「配列内の最大値・最小値の取得」を高速 (できれば定数時間) で行えるデータ構造が欲しい。これを実現するには、 C++ ではsetが適当だろう。

setでは、ポインタを用いることで配列内の最大値・最小値の取得を $O(1)$ で行える。

ここまで分かれば、あとは $a$ を全探索していけばよい。具体的なアルゴリズムはコード例の通り。

コード

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

template <typename T>
inline bool chmin(T &a, T b) { return ((a > b) ? (a = b, true) : (false)); }

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

int main()
{
    int N, K;
    cin >> N >> K;
    vector<int> v(N);
    rep(i, 0, N)
    {
        int P;
        cin >> P;
        P--;
        v[P] = i;
    }

    set<int> st;
    rep(i, 0, K) st.insert(v[i]);

    int ans = *st.rbegin() - *st.begin();
    rep(i, 0, N - K)
    {
        st.erase(v[i]);
        st.insert(v[i + K]);
        chmin(ans, *st.rbegin() - *st.begin());
    }

    cout << ans << endl;
}

atcoder.jp

実装時間: 40分

別解 : Segment Tree を利用する

上述の考察内では、「ある区間内の最大値・最小値」を求めることが最終的な目標となった。区間に関するクエリを処理するのに有効なデータ構造として、 Segment Tree というものがある。このデータ構造は C++ の標準ライブラリにはないのだが、実は AtCoder Library には用意されている。

github.com

別解として、これを使ったコードも載せておく。アルゴリズムとしては、最大値と最小値を求める用の Segment Tree をそれぞれ用意しておき、全探索によって両者の差の最小値を求めるというものである。

#include <bits/stdc++.h>
#include <atcoder/all>

using namespace std;
using namespace atcoder;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

constexpr auto INF = 1e+9;

template <typename T>
inline bool chmin(T &a, T b) { return ((a > b) ? (a = b, true) : (false)); }

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

int op_min(int a, int b) { return min(a, b); }
int e_min() { return INF; }

int op_max(int a, int b) { return max(a, b); }
int e_max() { return 0; }

int main()
{
    int N, K;
    cin >> N >> K;
    segtree<int, op_min, e_min> segMin(N);
    segtree<int, op_max, e_max> segMax(N);
    rep(i, 0, N)
    {
        int P;
        cin >> P;
        segMin.set(P - 1, i);
        segMax.set(P - 1, i);
    }

    int ans = INF;
    rep(i, 0, N - K + 1)
    {
        int min = segMin.prod(i, i + K);
        int max = segMax.prod(i, i + K);
        chmin(ans, max - min);
    }
    cout << ans << endl;
}