NoiminのNoise

競技プログラミング (多め) とWeb (たまに) ,自然言語処理 (ブログではまだ)。数式の書き方を一気に KaTeX に変えようとして記事を全削除してインポートし直すなどしたので,過去にブックマークされた記事は URL が変わってしまっている可能性があります…….

AtCoder Regular Contest 051 D - 長方形

問題文: D - 長方形

公式解説: http://arc051.contest.atcoder.jp/data/arc/051/editorial.pdf

問題概要

長さ $W (\leq 2000)$ の整数列 $ a_1, a_2, \cdots, a_W (a_i \leq 10^5)$ と,長さ $ H (\leq 2000)$ の整数列 $b_1, b_2, \cdots, b_H (b_i \leq 10^5)$ が与えられる.

左から i 列目,上から j 列目のマス目に $a_i + b_j$ が書き込まれているような $H \times W$ のマス目を考えるとき,次の内容の $Q (\leq 2000)$ 個のクエリに答えよ.

  • A,B が与えられるので,左から A 列目以内かつ上から B 行目以内のマス目たちの中から,選んだ部分が長方形になるようにマス目を選んだ時の,選んだマス目の値の総和の最大値.

解法概要

$ \mathit{sa}_{i, w} = \sum_{k=i-w+1}^{i} a_k $,$ \mathit{sb}_{j, h} = \sum_{k=j-h+1}^{j} b_k $ とおくと,クエリの答えは次のように表せる:

$\max_{1 \leq w \leq A, 1 \leq h \leq B, w \leq i \leq W, h \leq j \leq H } h \times \mathit{sa}_{i, w} + w \times \mathit{sb}_{j, h}$ .

これを愚直に計算すると1クエリあたりの時間計算量が $O(H^2 W^2)$ で明らかに間に合わない.

まず,w, h が固定されているならば $\max_{1 \leq k \leq i} \mathit{sa}_{k, w}$ は $ \max_{1\leq k \leq i-1} \mathit{sa}_{k, w}$ と $ \mathit{sa}_{i, w}$ から求めることができる.仮に $\mathit{sa}_{i, w}$ を全ての i, w について計算していても時間・空間計算量は $O(W^2)$ なので大丈夫そう.なので,$\mathit{maxSa}_{i, w} = \max_{1 \leq k \leq i} \mathit{sa}_{i, w}$ をあらかじめ用意しておく.同様に,$\mathit{maxSb}_{j, h} = \max_{1 \leq k \leq j} \mathit{sb}_{j, h}$ も用意しておく.

すると,クエリの答えは次のように変形できる:

$\max_{1 \leq w \leq A, 1 \leq h \leq B } h \times \mathit{maxSa}_{A, w} + w \times \mathit{maxSb}_{B, h}$.

上記のような一次式の集合に対して max を求めるテクとして,Convex Hull Trick がある.Convex Hull Trick は px+q という式で表せる直線の集合について, x にある値を与えたときに最大値を取る直線はどれか,または最大値そのものを高速に求めることができるテクニックである.与えられる x の数,言い換えればクエリの数が X 個で,直線集合に含まれる直線の数が L 本であるならば,クエリごとに愚直に計算すれば時間計算量 $O(XL)$ となるところを $O((X+L) \log L)$で処理することが可能になる.Convex Hull Trick そのものについての知識や実装方法については,satanic0258 さんのブログの記事が非常にわかりやすかった.おすすめ.

さて,Convex Hull Trick を適用できるようにもう少し式変形を行う. w を最初に固定してから h について最大化し, それぞれの w の結果の中から最大値を選ぶとすると,次のような変形ができる:

$ \max_{1 \leq w \leq A } \left\{ w \max_{1 \leq h \leq B} \left( h \times \frac{\mathit{maxSa}_{A, w}}{w} + \mathit{maxSb}_{B, h} \right) \right\} $ .

先ほどの Convex Hull Trick の説明で直線の式を px+q と表現したが,$ p = \frac{\mathit{maxSa}_{A, w}}{w} $ ,$q = \mathit{maxSb}_{B, h}$ とした B本の直線を考えると, $\max_{1 \leq h \leq B} \left( h \times \frac{\mathit{maxSa}_{A, w}}{w} + \mathit{maxSb}_{B, h} \right)$ の部分の計算が Convex Hull Trick で高速化できる (当該部分のみの時間計算量は $O(\log B)$).あとはその値に w をかけた値についての最大値をとれば,それがそのクエリの答えになる.

以上のような計算をすると,1クエリごとの時間計算量は $O((B+A) \log B) \subseteq O((H+W) \log H)$ となり,クエリの数が上限いっぱいいっぱいでも間に合う.

ソースコード

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

using ll =  long long;
using Pll = pair<ll, ll>;

constexpr ll INF = 1LL << 60;
constexpr double EPS = 1e-5;

class ConvexHullTrick {
    public:
    int n = 0;
    vector<Pll> lines;
    
    ConvexHullTrick() {}

    bool check(Pll line1, Pll line2, Pll line3) {
        return (line3.second-line2.second)*(line1.first-line2.first) > (line2.second-line1.second)*(line2.first-line3.first);
    }

    // 直線の傾きの降順に直線が与えられる前提
    void add(ll grad, ll intercept) {
        while(n >= 2 && !check(lines[n-2], lines[n-1], {grad, intercept})) {
            lines.pop_back();
            --n;
        }
        lines.emplace_back(grad, intercept);
        ++n;
    }

    double f(ll i, double x) {
        return lines[i].first*x + lines[i].second;
    }

    double get_min(double x) {
        int ok = 0, ng = n;
        while(ng-ok > 1) {
            int mid = (ok+ng)/2;
            if(mid && f(mid-1, x) >= f(mid, x)) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        return f(ok, x);
    }
};

int main() {
    int W, H;
    cin >> W >> H;
    vector<ll> a(W+1, 0), b(H+1, 0), sa(W+1, 0), sb(H+1, 0);
    vector<vector<ll>> max_sa(W+1, vector<ll>(W+1, -INF)), max_sb(H+1, vector<ll>(H+1, -INF));
    for(int i=1;i<=W;++i) {
        cin >> a[i];
        sa[i] = sa[i-1] + a[i];
    }
    for(int i=1;i<=H;++i) {
        cin >> b[i];
        sb[i] = sb[i-1] + b[i];
    }
    for(int w=1;w<=W;++w) {
        for(int k=w;k<=W;++k) {
            max_sa[k][w] = max(max_sa[k-1][w], sa[k] - sa[k-w]);
        }
    }
    for(int h=1;h<=H;++h) {
        for(int k=h;k<=H;++k) {
            max_sb[k][h] = max(max_sb[k-1][h], sb[k] - sb[k-h]);
        }
    }

    int Q, A, B;
    cin >> Q;
    while(Q--) {
        cin >> A >> B;

        ConvexHullTrick cht = ConvexHullTrick();
        for(int h=1;h<=B;++h) {
            cht.add(-h, -max_sb[B][h]);
        }

        ll ans = a[A] + b[B];
        for(int w=1;w<=A;++w) {
            double tmp_ans = double(w) * -cht.get_min(double(max_sa[A][w]) / w);
            ans = max(
                ans, 
                ll(tmp_ans + (tmp_ans > 0 ? 1: -1) * EPS)
            );
        }

        cout << ans << endl;
    }
}

所感

Convex Hull Trick,知ってはいたし前にこどふぉの問題を解いたりもしたけど,実戦投入にはまだまだ精進が必要そう.