yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

C - Theofanis' Nightmare / Codeforces Round 912 (Div. 2)

CF1903 Codeforces C++ Div.2 Codeforces Round 912 Diff 1400 貪欲法 競技プログラミング

問題

n 個の要素からなる整数列 a が与えられる。

これをいくつかの空でない部分列に分割する。 i 番目の部分列の要素の合計を \mathrm{sum}_{i} とし、部分列の個数を k 個とするとき、 \displaystyle \sum_{i = 1}^{k} i \cdot \mathrm{sum}_{i} の最大値を求めよ。

例

例えば、 a = [1, -3, 7, -6, 2, 5 ] とし、この数列 a を [ 1 ], [-3, 7], [-6, 2], [5] と分割したとき、 \begin{align} \sum_{i = 1}^{k} i \cdot \mathrm{sum}_{i} = 1 \cdot (1) + 2 \cdot (-3 + 7) + 3 \cdot (-6 + 2) + 4 \cdot (5) = 17 \end{align} と計算される。

入力

まず最初の1行目に、テストケースの個数を表す整数 t が与えられる。

その後、 t 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

  • 最初の1行目には、整数 n が与えられる。
  • 次の1行には、 n 個の整数 {a}_{i} が順に与えられる。

条件

  • 1 \leq t \leq {10}^{4}
  • 1 \leq n \leq {10}^{5}
    • すべてのテストケースでの n の値の合計は 2 \cdot {10}^{5} を超えない。
  • - {10}^{8} \leq {a}_{i} \leq {10}^{8}

出力

各テストケースにおける最大値を1行ずつ出力すること。

解法

数列 a の要素すべてが 0 以上であるとき、合計値として最大となるのは \begin{align} 1 \cdot {a}_{1} + 2 \cdot {a}_{2} + \cdots + (n - 1) \cdot {a}_{n - 1} + n \cdot {a}_{n} \end{align} のときです。 すなわち、数列 a を [{a}_{1}], [{a}_{2}], \cdots, [{a}_{n}] と各要素を1つの部分列として扱うときに、求める値が最大となります。

これは、例えば数列 a を [{a}_{1}], [{a}_{2}], \cdots, [{a}_{i - 1}], [{a}_{i}, {a}_{i+1}], [{a}_{i + 2}], \cdots, [{a}_{n}] のように、 {a}_{i} と {a}_{i + 1} を1つの部分列として扱ったときの値は \begin{align} 1 \cdot {a}_{1} + \cdots + (i - 1) \cdot {a}_{i - 1} + i \cdot ({a}_{i} + {a}_{i + 1}) + (i + 1) \cdot {a}_{i + 2} + \cdots + (n - 1) \cdot {a}_{n} \end{align} となることから、すべての要素を1つの部分列として扱うときの値と差を取ると、 \begin{align} & \ (1 \cdot {a}_{1} + \cdots + n \cdot {a}_{n}) \\ - & \ \{ 1 \cdot {a}_{1} + \cdots + (i - 1) \cdot {a}_{i - 1} + i \cdot ({a}_{i} + {a}_{i + 1}) + (i + 1) \cdot {a}_{i + 2} + \cdots + (n - 1) \cdot {a}_{n} \} \\ = & \ {a}_{i + 1} + {a}_{i + 2} + \cdots + {a}_{n} \end{align} となり、これは0以上の値となります。

従って、すべての i に対して {a}_{i} \geq 0 が成立するときは、各要素を1つの部分列とするときが最大となります。 すなわち、なるべく要素数を増やし、掛けられる i の値を大きくすることによって、最大値が得られるということが言えます。

一方で、今回の問題では {a}_{i} \lt 0 である場合も存在します。

ここで、数列 a の i + 1 番目以降の要素について、分け方が決定していると仮定します。 （ a の 1 番目から i 番目の要素については、すべて要素数1の部分列として分割されているものとします。） このときの求める値は \begin{align} 1 \cdot {a}_{1} + 2 \cdot {a}_{2} + \cdots + i \cdot {a}_{i} + (i + 1) \cdot ({a}_{i + 1} + \cdots) + \cdots \end{align} となります。

このとき、 {a}_{i} をどの部分列に所属させるかということについては、

  1. [ {a}_{i} ] として独立させる
  2. 隣の部分列 [ {a}_{i + 1}, \cdots ] と統合して、 [ {a}_{i}, {a}_{i + 1}, \cdots ] とする

の2通りの方法が考えられます。

このとき1つ目の方法については、分け方について変化はないので、値についても変わりはありません。

一方で2つ目の方法については、 i + 1 番目以降の部分列の番号が1つずつ変わります。 （ {a}_{i + 1} の所属する部分列の番号が i + 1 から i となります。）

そのため、元の値と比較すると、 \begin{align} & \ \{ 1 \cdot {a}_{1} + \cdots + (i - 1) \cdot {a}_{i - 1} + i \cdot {a}_{i} + (i + 1) \cdot ( {a}_{i + 1} + \cdots ) + \cdots \} \\ - & \ \{ 1 \cdot {a}_{1} + \cdots + (i - 1) \cdot {a}_{i - 1} + i \cdot ({a}_{i} + {a}_{i + 1} + \cdots ) + \cdots \} \\ = & \ {a}_{i + 1} + \cdots + {a}_{n} \end{align} となります。

以上から、 {a}_{i + 1} + \cdots + {a}_{n} の値が0以上である場合は元の値の方が大きくなり、そうでない場合は新しい値の場合の方が大きくなります。 従って、この値が0以上のときは [{a}_{i}] として独立させ、そうでないときは [{a}_{i}, {a}_{i + 1}, \cdots ] として隣の部分列に統合することにより、求める値をより大きくしていくことができます。

これを i = n - 1, n - 2, \cdots, 1 の順で行っていくことにより、求める値の最大値を求めることができます。

ここで、値 \mathrm{total}_{i} = {a}_{1} + \cdots + {a}_{i} と定義すると、 \begin{align} {a}_{i + 1} + \cdots + {a}_{n} = \mathrm{total}_{n} - \mathrm{total}_{i} \end{align} が成立するため、 {a}_{i + 1} + \cdots + {a}_{n} の値を O(1) で求めることができます。

従って、この値 \mathrm{total} の準備を予め行うことにより、各テストケースに対して O(n) の計算量で答えを出力することができます。

ソースコード

与えられた数列に対して、最大値を返す関数 solve() を以下のように実装しました。

int n;
ll a[int(1e5 + 5)], total[int(1e5 + 5)];

ll solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  // 数列 total の準備
  for (int i = 1; i <= n; i++) {
    total[i] = total[i - 1] + a[i];
  }

  ll ans = 0; // 答えとなる値
  for (ll i = 1; i <= n; i++) {
    ans += i * a[i]; // まずは、1要素ずつが部分列の場合の値を求める
  }

  for (int i = n - 1; i >= 1; i--) {
    ans = max(ans, ans - (total[n] - total[i])); // a[i] を今のままにする場合と次の部分列に入れる場合との値を求め、大きい方を ans として更新する
  }

  return ans;
}