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

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

G - 区間和 / 第13回 アルゴリズム実技検定 (PAST13)

PAST13 AtCoder C++ アルゴリズム実技検定 PAST DP 区間和 競技プログラミング

問題

長さ n の整数列 a = [ {a}_{1}, {a}_{2}, \cdots , {a}_{n} ] が与えられる。

1 \leq l \leq r \leq n を満たす整数組 (l, r) を任意に選ぶときの、 {a}_{l} + {a}_{l+1} + \cdots + {a}_{r} の値としてありえる最大値を求めよ。

入力

まず最初の1行目に整数 n が与えられる。

その次の1行に n 個の整数 {a}_{1}, {a}_{2}, \cdots , {a}_{n} が順に与えられる。

条件

  • 1 \leq n \leq 2 \times {10}^{5}
  • - {10}^{9} \leq {a}_{i} \leq {10}^{9}
  • 入力はすべて整数

出力

最大値を1行で出力すること。

解法

まず単純な方法として、 (l, r) となる値を 1 から n まで順に探索し、和を求めていくという方法が着想できます。

この方法では、ループを行うだけでも O({n}^{2}) の計算量が必要になり、 n は最大で 2\times{10}^{5} であるので、実行時間制限には間に合いません。 そのため、別の方法を考える必要があります。

そこで、以下のような値 \mathrm{dp1}_{i}, \mathrm{dp2}_{i} を考えます。

  • \mathrm{dp1}_{i} : r = i としたときの {a}_{l} + {a}_{l+1} + \cdots + {a}_{i} の最大値
  • \mathrm{dp2}_{i} : l = i としたときの {a}_{i} + {a}_{i+1} + \cdots + {a}_{r} の最大値

これらの値を用いると、 \begin{align} \mathrm{dp 1}_{i} + \mathrm{dp 2} _{i} - {a}_{i} \end{align} が、 {a}_{i} を含んでいるときの {a}_{l} + \cdots + {a}_{r} の最大値となると言えます。 （ここで、 {a}_{i} を引いているのは、 \mathrm{dp1}_{i} と \mathrm{dp2}_{i} で {a}_{i} が重複しているためです。）

従って、これを各 i について探索していくことにより、 O(n) の計算量で最大値を求めることができると言えます。

あとは、 \mathrm{dp1}_{i}, \mathrm{dp 2}_{i} のそれぞれを求めることを考えます。

\mathrm{dp1}_{i} は、 {a}_{i} を右端としたときの和の最大値です。 そのため、 \mathrm{dp1}_{ i + 1 } の値として考えられるのは、

  • i までの最大値に {a}_{i+1} を足した、 \mathrm{dp1}_{i} + {a}_{i + 1}
  • {a}_{i + 1} 自体

の2通りのいずれかとなります。 従って、 \begin{align} \mathrm{dp 1}_{ i + 1 } = \max ( \mathrm{dp 1}_{ i } + {a}_{i + 1}, {a}_{i + 1} ) \end{align} が成立します。 これを、 i = 1, 2, \cdots , n の順で計算していくことにより、 \mathrm{dp 1}_{i} の準備ができます。

\mathrm{dp2}_{i} も {a}_{i} を左端としたときの和の最大値であるので、同様の考え方から \mathrm{dp2}_{i - 1} の値として考えられるのは、

  • i からの最大値に {a}_{i-1} を足した、 \mathrm{dp2}_{i} + {a}_{i - 1}
  • {a}_{i - 1} 自体

の2通りのいずれかとなります。 従って、 \begin{align} \mathrm{dp 2}_{i - 1} = \max ( \mathrm{dp 2}_{i} + {a}_{i - 1}, {a}_{i - 1} ) \end{algin} が成立します。 これを、 i = n, n - 1, \cdots, 1 の順で計算していくことにより、 \mathrm{dp 2}_{i} の準備ができます。

以上で準備された \mathrm{dp1}_{i}, \mathrm{dp2}_{i} を用いて、 i = 1, 2, \cdots, n に対して \begin{align} \mathrm{dp 1}_{i} + \mathrm{dp 2} _{i} - {a}_{i} \end{align} の値を求めていき、その中の最大値が求めるものとなります。

ソースコード

main() 関数の中に、直接答えを出力する部分を実装しました。

int n;
ll a[int(2e5 + 5)];
ll dp1[int(2e5 + 5)], dp2[int(2e5 + 5)];

int main() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }

  // dp1 の準備
  for (int i = 1; i <= n; i++) {
    dp1[i] = max(dp1[i - 1] + a[i], a[i]);
  }
  // dp2 の準備
  for (int i = n; i >= 1; i--) {
    dp2[i] = max(dp2[i + 1] + a[i], a[i]);
  }

  ll ans = a[1]; // 答えとなる最大値 ans （初期値に 0 を指定すると、要素が全て負のときなどに対応できない）
  for (int i = 1; i <= n; i++) {
    ans = max(ans, dp1[i] + dp2[i] - a[i]); // 最大値の更新
  }
  cout << ans << endl; // 答えの出力

  return 0;
}

感想

問題設定がシンプルで、それでいて計算量削減のためにどのようなことができるのかを考えられる問題なので、すごくいいなと思いました。