第二回 アルゴリズム実技検定 L - 辞書順最⼩ 解説

問題

問題文は本家サイトにあります：第二回 アルゴリズム実技検定 L - 辞書順最⼩

問題概要

長さ $N$ の整数列 $A_1, A_2, \cdots, A_N$ が与えられます。ここから $K$ 個の要素を選び、相対的な順序を保ったまま並べて新しい数列を作ります。ただし、 $A_i, A_j$ をともに選べるのは $|i-j|\geqq D$ であるときに限ります。

こうしてできる数列のうち、辞書順で最小のものを答えてください。そのような数列が作れない場合は $-1$ を出力してください。

制約

$2\leqq N \leqq 2\times 10^5$
$1\leqq D, K \leqq N$
$0\leqq A_i \leqq 10^9$

解説

例えば、次のような数列に対して考えてみましょう。 $K=3$, $D=2$ とします。つまり、 $|i-j|\geqq 2$ を満たすように要素3つを選ぶ、ということです。

7 2 3 5 4 4 1 6

辞書順で最小にしたいので、新しい列については一番左から考えていくのがいいでしょう。まずは、1を選びたいところです。しかし、これを選ぶと、3個の数字が選べません。では、2を選ぶのはどうでしょうか。この場合、後で5と6を選ぶといったことができるので、2を選ぶのは大丈夫そうです。7や3も選べますが、選んだ時点で負けが確定なので、捨てていい選択肢です。

次は、やはり1を選びたいですが、残りの1個が選べないのでダメですね。また、3を選ぶこともできません。2文字目と3文字目は $D=2$ 以上離れていないからです。結局、4を選ぶのがいいことがわかります。

4は2つありますが、もし右の方を選んでしまうと、残りは6を選ぶしかないですね。一方、左の4を選べば、1も選ぶことができます。こうして、新しい数列として、

2 4 1

と選べばいいことがわかります。

この例から、次のように決めていくのがよさそうです。まず、新しい数列は左から順番に決めていくことにします。左の要素が小さければ、辞書順で必ず小さくなるからです。それぞれの要素は、今選べる範囲で一番小さな数字を選びます。同じ要素があれば、その中で一番左の要素を選びます。そのほうが、今後選べる選択肢が増えるからです。これを繰り返していきます。

さて、これをやるためには、今選べる範囲をどう管理するかを考えなくてはいけません。これは、左側と右側をそれぞれ管理すればいいでしょう。

左は、「1つ前に $i$ 文字目を選んだ場合、次は $i+D$ 文字目以降」から選べばいいですね。右側は、選べるギリギリの範囲を考えればいいです。先ほどの例を使ってみます。

7 2 3 5 4 4 1 6

新しい数列の1つ目の要素を選ぶ際、1つ目として選べるギリギリの右側の要素はどこでしょうか。それは3つ目に6を選び、2つ目に右側の4を選んだときです。このとき、1文字目として5を選ぶことができます。つまり、4番目の要素です。同様に考えて、新しい数列の2つ目の要素は6番目まで、3つ目の要素は8番目の要素まで選ぶことができます。

これを式で書くには、右から考えたほうがいいでしょう。新しい数列の最後の要素は $N$ 番目の要素まで選べます。その1つ前の要素は、 $N-D$ 番目の要素まで選べます。以下、 $D$ ずつさかのぼっていけばいいので、新しい数列の1つ目の要素は、元の数列の $N-(K-1)D$ 番目まで選べることがわかります。

選べる範囲の中から、値が小さいものを選びます。値が同じなら位置が左にあるものを選びます。これを効率よく行うには、優先度付きキュー priority_queue を使うのがいいです。数列の値と位置とをペアにし、小さい順からとりだせるようにします。

以上のことを踏まえると、C++では以下のように書けます。コードの下に説明が続きます。

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

int main() {

  ll N, K, D; cin >> N >> K >> D;
  vector<ll> A(N); for (ll i = 0; i < N; i++) cin >> A[i];

  priority_queue<pll, vector<pll>, greater<pll>> q;
  ll l = 0, r = N - (K - 1) * D;
  for (ll i = 0; i < r; i++) {
    q.push(make_pair(A[i], i));
  }
  if (q.empty()) {
    cout << "-1\\n";
    return 0;
  }

  queue<ll> ans;
  for (ll i = 0; i < K; i++) {
    while (true) {
      pll p = q.top(); q.pop();
      if (l <= p.second) {
        ans.push(p.first);
        if (r == N) break;
        for (ll j = r; j < r + D; j++) {
          q.push(make_pair(A[j], j));
        }
        l = p.second + D;
        r += D;
        break;
      }
    }
  }

  for (ll i = 0; i < K; i++) {
    ll p = ans.front(); ans.pop();
    cout << p << ' ';
  }
  return 0;
}

要素を選べる範囲の左端と右端を l, r で管理しています。初期値では l = 0 で、 r = N - (K - 1) * D です（0-indexed）。

次に、元の数列の要素と位置をペアにして、 priority_queue q に入れていきます。選べる範囲だけを対象にしたいので、 r 未満だけ入れます。ここで、 q の中身が空であれば、新しい数列が作れないということなので -1 を出力して終わります。

続いて、新しい数列を作っていく作業です。 q から要素を出します。これが、今選べる範囲のものなら選ぶことにします。そして、新しく選べるようになった要素を追加し、 l, r を更新します。 l は、選んだ要素の場所から +D し、 r は +D します。あとはこれの繰り返しです。こうして、選べる範囲を更新しながら、値を選んでいくことができます。

選んだものを queue ans に入れておき、最後に出力しておしまいです。if (r == N) break; は、 K が 1 のときに、値の追加をおこわ内容にするための処理です。

先ほどの例を使って、どのように更新されていくか、具体的に見ていきます。

7 2 3 5 4 4 1 6

この数列に対し、まずは次の範囲を考えます。

7 2 3 5 4 4 1 6

この中で一番小さいのは 2 なので、2を選びます。 l は2の場所から2つ右、 r も2つ右にずれて、次の要素は次の範囲から選びます。

7 2 3 5 4 4 1 6

この中で一番小さく、左にあるものは、左側の4です。 l, r は更新されて、次の範囲はこうです。

7 2 3 5 4 4 1 6

これで1を選びます。上のコードでは、こういった流れで新しい数列を作っています。