Mister雑記

競プロ記事がメインになりますが、内容は気分次第です。

ARC074-D 「3N Numbers」 (500)

数列系の典型が詰められた良問。スラスラ実装できるようにしておきたい。

問題リンク

概要

長さ 3Nの数列 aが与えられる。ここから順番を変えず N個の要素を取り除いて長さ 2Nの数列 a'を作る。このとき、 (a'の前半 N要素の総和 )-(a'の後半 N要素の総和)の最大値を求めよ。

制約

  •  1 \leq N \leq 10^5
  •  1 \leq a_i \leq 10^9

解説

解の求め方

長さ 3Nの数列から長さ 2Nの数列を作り、その前半 N要素を赤に、後半 N要素を青に塗ってみる。このとき、元の数列において赤い要素は全て青い要素より左側にある

当然といえば当然な性質だが、これがとても重要。答えとなる塗り分けにおいても、ある一箇所で数列 aを切断することでその前半には赤い要素だけ、後半には青い要素だけが含まれていることになる。

最善かは分からないが、試しに切断箇所を1つ固定する。このとき、以下のように前半は大きい要素から N個、後半は小さい要素から N個塗るのが最善となるのは分かるだろう。

f:id:misteer:20180917101625p:plain

まとめると、全ての考えうる切断箇所について上のように適切に塗り分けをし、値を求めてその最大値を取れば答えとなるわけだ。

計算量の削減

しかし切断するごとに毎回ソートして、大きい/小さい方から N個選んでいるようでは計算量が O(N^2)でTLEとなる。

そこで役に立つのが、優先度付きキューと呼ばれるデータ構造である。C++ではpriority_queueという名前が付いている。

このデータ構造は、何らかの「評価」において一番評価の高いものに O(1)でアクセスでき、さらにO(1) O(\log N)で削除もできるという性質を持つ。「評価」が「大きさ」なら、一番大きい要素を調べたり削除するのが超高速にできるということだ。

これを使えば、以下のようにして赤に塗り分けられた値の総和を順番に調べることができる。

  • 「評価」が「小ささ」のpriority_queueを用意する。
  •  a_1 \sim a_Nをpriority_queueに突っ込み、その総和を red_Nとしておく。
  • 以降、以下のようにして i \in (N, 2N]について red_iを求める。
    •  red_i = red_{i - 1} + a_iとし、priority_queueに a_iを突っ込む。
    • priority_queueから一番小さい要素を取り出し、 red_iから引く。
    • その要素をpriority_queueから削除する。

こうすることで、priority_queueには常に大きい方から N個の要素が入っていることになる。青も大小と追加する順番を逆にするだけで同様に実装すればいい。

実装例

全体として1-indexedであることに注意。

#include <algorithm>
#include <iostream>
#include <queue>

using namespace std;

using ll = long long;
template <typename T>
using PQ = priority_queue<T>;
template <typename T>
using GPQ = priority_queue<T, vector<T>, greater<T>>;

const ll INF = 1LL << 50;

int main() {
    ll N;
    cin >> N;
    ll a[N * 3 + 1];
    for (int i = 1; i <= N * 3; ++i) {
        cin >> a[i];
    }

    ll red[N * 2 + 1];
    red[0] = 0;
    // red[i] = 前半i個からN個選んだ総和の最大値

    GPQ<ll> rque;
    // 大きい方からN個を保持

    for (int i = 1; i <= N; ++i) {
        red[i] = red[i - 1] + a[i];
        rque.push(a[i]);
    }

    // ここから小さいのが切り捨てられる
    for (int i = N + 1; i <= N * 2; ++i) {
        red[i] = red[i - 1] + a[i];
        rque.push(a[i]);
        red[i] -= rque.top();
        rque.pop();
    }

    // aをひっくり返して同じことをする
    reverse(a + 1, a + N * 3 + 1);

    ll blue[N * 2 + 1];
    blue[0] = 0;
    // blue[i] = 後半i個からN個選んだ総和の最小値

    PQ<ll> bque;
    // 小さい方からN個を保持

    for (int i = 1; i <= N; ++i) {
        blue[i] = blue[i - 1] + a[i];
        bque.push(a[i]);
    }

    for (int i = N + 1; i <= N * 2; ++i) {
        blue[i] = blue[i - 1] + a[i];
        bque.push(a[i]);
        blue[i] -= bque.top();
        bque.pop();
    }

    ll ans = -INF;
    for (int i = N; i <= N * 2; ++i) {
        // 前半i個、後半3N - i個の差
        ans = max(ans, red[i] - blue[N * 3 - i]);
    }

    cout << ans << endl;
    return 0;
}

priority_queueはDijkstra法などでも使われるデータ構造なので、そっちで慣れている人も多いだろう。なんにせよ、使う機会を増やして慣れておくに越したことはない。

感想

  • この「どこかで分割してhogehogeする」ってのは数列系では相当強力な手法なので、すぐに浮かぶようにしておきたい。
  • 初めて解いた5ヶ月前よりは遥かにスムーズに実装できるようになっていた(それはそう)。
  • 実は5本の指に入るくらい好きな問題。

追記

priority_queueの最悪計算量は

  • 挿入が O(\log N)
  • 先頭の削除が O(\log N)
  • 先頭要素のアクセスが O(1)

です。本文中に誤りがあったので訂正しておきます。