Skyscrapers (hard ver)

codeforces.com

問題概要

簡単に言えば

1. ある配列が与えられる
2. ある要素1つを選ぶ
3. その要素より左は広義単調増加、右は広義単調増加ように配列の各要素の値を減らす
4. 操作後の配列のうち総和が最大となるときの配列を出力せよ。

C1と同様の問題で、制約のみが異なる。

codeforces.com

解法

C1と同様に、ある要素を選んで実際に配列を構築するやりかたは計算量がO(n²)かかる。今回は制約がn <= 5e5なので通らない。

ある位置を選んだ後、実際に配列を構築するのはO(n)で可能。つまり、今回ネックとなっているのは「どの位置を選ぶべきか？」の判定部分である。

第一感としては累積和でいけそうである。そして実際にそれでいける。

まず、基本的には累積和をとる。ただし、左側に存在する「今見ている要素より大きい要素」は今見ている要素に置き換えて和を求める。この作業はO(n)かかりそうだが(実際最悪かかる)、うまくやると償却計算量がO(1)となる。やり方としてはmapで各要素のカウントを取り、priority queueで出てきた要素の種類を保持すると計算量を落とせる。*1

同様の方法を右側からも行う。こうしてO(n)で左から、右からの累積和を求めることができた。

あとは簡単に答えを求めることが出来る。

1. 左からの累積和、右からの累積和の2つを取る(O(n))
2. それらを足すことで、ある位置を選んだときの配列の総和を求める(O(1))
3. 総和が最大となる位置を調べる(O(n))

全体としてO(n)となる。

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <map>

using namespace std;
using ll = long long;

#define in(v) v; cin >> v;
#define all(f,c,...) (([&](decltype((c)) cccc) { return (f)(begin(cccc), end(cccc), ## __VA_ARGS__); })(c))
#define _overload3(_1,_2,_3,name,...) name
#define _rep(i,n) for(int i=0;i<(n);++i)
#define repi(i,a,b) for(int i=(a);i<(b);++i)
#define rep(...) _overload3(__VA_ARGS__,repi,_rep,)(__VA_ARGS__)
#define rrep(i,n) for(int i=(n);i>=0;--i)

template<class T>bool chmax(T &a,const T &b){if(b>a){a=b;return 1;}return 0;}
template<class T>bool chmin(T &a,const T &b){if(b<a){a=b;return 1;}return 0;}

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


  priority_queue<int> q;
  map<int, int> m, m2;
  ll sum = 0;
  vector<ll> sum_l(N+1), sum_r(N+1);
  rep(i, N) {
    // 左から累積和をとっていく
    // ただし、左に存在する要素はすべて今見ている要素以下になるよう調整されているものとする
    sum += A[i];
    ++m[A[i]];
    q.push(A[i]);
    while (!q.empty() && q.top() > A[i]) {
      // このループは最悪計算量はO(n)だが、償却計算量がO(1)
      int p = q.top(); q.pop();
      sum += (ll)m[p] * (A[i] - p);
      m[A[i]] += m[p];
      m[p] = 0;
    }
    sum_l[i+1] = sum;
  }
  sum = 0;
  rrep(i, N-1) {
    // 以下は上のコードのコピペ。右からの累積和を取っているだけ
    sum += A[i];
    ++m2[A[i]];
    q.push(A[i]);
    while (!q.empty() && q.top() > A[i]) {
      int p = q.top(); q.pop();
      sum += (ll)m2[p] * (A[i] - p);
      m2[A[i]] += m2[p];
      m2[p] = 0;
    }
    sum_r[i] = sum;
  }
  int max_i = 0;
  rep(i, N)
    if (chmax(sum, sum_l[i] + sum_r[i])) max_i = i;


  rrep(i, max_i-1) chmin(A[i], A[i+1]);
  rep(i, max_i, N-1) chmin(A[i+1], A[i]);
  rep(i, N) {
    if (i) cout << " ";
    cout << A[i];
  }
  cout << endl;

  return 0;
}

*1:うまく説明できないのだが、このやり方なら最悪計算量はO(n)となる。まあ5, 4, 3, 2, 1のような降順の配列に対する操作を考えていただければ、直感としてはわかっていただけるのではないだろうか。