🐍

ABC 438 D 問題:後方累積最大値(末尾累積 max)の扱い方

に公開

はじめに

本記事では 後方累積最大値 (末尾累積 max)の構成法と使い方の型を説明します。

まずは累積和に関する用語の整理と定義を行い、例題で後方累積最大値の使い方を見ていきます。続いてこのデータ構造が持つ性質を確認し、最終的に ABC 438 D 問題への応用に繋げるという流れで進めます。
初期値の設定や添字の扱いで迷わずに扱えるようになることが本記事の目標です。

なお ABC 438 D 問題の詳細な解説は本記事では行いませんので、必要な場合は公式解説などと併せてお読みください。

累積和と後方累積和

用語整理のため、まずは累積和・後方累積和について軽く触れます。
以降、記事内の添字はすべて 0-indexed です。

累積和(前方累積和)はリストの先頭から順に和を取ったデータ構造です。
また、問題によっては末尾から順に和を累積したデータ構造が必要になる場合があります。こうしたデータ構造は末尾累積和や後方累積和と呼ばれますが、本記事では 後方累積和 と統一します。

後述の累積最大値との対比のため、数列 A に対する典型的な累積和 P_{\rm{sum}} および 後方累積和 S_{\rm{sum}} の形を以下に示します。

\begin{align*} P_{\rm{sum}} & = \bigl( 0,\ A_0,\ A_0 + A_1,\ \dots,\ A_0 + \cdots + A_{N - 1} \bigr) \\ S_{\rm{sum}} & = \bigl( A_0 + \cdots + A_{N - 1},\ A_1 + \cdots + A_{N - 1},\ \dots,\ A_{N - 1},\ 0 \bigr) \\ \end{align*}

先頭(末尾)に 0 を挿入していますが、これは半開区間で扱いやすくするためです。

累積最大値と後方累積最大値

続いて累積最大値の定義を述べ、以降は本題である後方累積最大値に焦点を当てていきます。

累積和の考え方を一般化し、先頭から順に \max を取って累積したデータ構造を 累積最大値 P_{\rm{max}} と呼びます。
さらにこれを末尾から順に累積したものが 後方累積最大値 S_{\rm{max}} です。

累積和で 0 を挿入したように、累積最大値では -\infty (または十分小さい値)を挿入します。

\begin{align*} P_{\rm{max}} & = \bigl( -\infty,\ \max(A_0),\ \max(A_0, A_1),\ \dots,\ \max(A_0, \dots, A_{N - 1}) \bigr) \\ S_{\rm{max}} & = \bigl( \max(A_0, \dots, A_{N - 1}),\ \max(A_1, \dots, A_{N - 1}),\ \dots,\ \max(A_{N - 1}),\ -\infty \bigr) \\ \end{align*}

この形式で定義することで、以下に示す性質が成り立つようになります。

後方累積最大値の性質

A を長さ N の数値のリストとして、後方累積最大値 S_{\rm{max}} の性質を列挙します。

  • S_{\rm{max}} の長さは N + 1 である。
  • S_{\rm{max}}[i]A_i 以降(i を含む) の値の最大値である。
  • S_{\rm{max}}[0] = \max(A) が成り立つ。
  • S_{\rm{max}}[N]-\infty (または十分小さい値)である。
  • i \lt j について常に S_{\rm{max}}[i] \ge S_{\rm{max}}[j] が成り立つ。

注意すべき点として、S_{\rm{max}}[i] を「i を含まない i より後方の値の最大値」のように定義してしまうと、これらの性質が成り立たなくなります。

では、次節でこれらの性質を用いた例題を見ていきます。

例題:各位置より後方にある要素の最大値をすべて求める

  • 長さ N の数列 A が与えられる。数列内の位置 0 \le i \lt N - 1 について、i より後方の最大値 \max_{i \lt j \lt N}(A_j) を求め、その値を i の昇順に出力せよ。
  • 制約は以下のとおり。
    • 2 \le N \le 10^6
    • -10^6 \le A_i \le 10^6 \quad (0 \le i \lt N)

この問題は後方累積最大値を前計算することにより計算量 \mathcal{O}(N) で解くことができます。

  1. まず後方累積最大値 S_{\rm{max}} を末尾から順に計算する。
  2. 次に 0 \le i \lt N - 1 に対して S_{\rm{max}}[i + 1] を出力する。

この例題に対する解法を Python で書くと以下のようになります。

n = int(input())
a = list(map(int, input().split()))

# 後方累積最大値の前計算
INF = 10**18
smax = [-INF] * (n + 1)
for i in range(n - 1, -1, -1):
    smax[i] = max(smax[i + 1], a[i])

# 解の出力
for i in range(n - 1):
    print(smax[i + 1])

最後に、このように「数列のある位置より後方の最大値を高速に求められる」ことを ABC 438 D 問題に応用した例を紹介します。

ABC 438 D 問題

問題全文:D - Tail of Snake

以下の内容は問題原文を 0-indexed に置き換えています。
長さ N の 3 つの数列 A, B, C が与えられるので、区切り位置を表す整数の組 (x, y) を適切に選んで以下に示す値 T の最大値を求める問題です。
ただし (x, y)1 \le x \lt y \lt N を満たす必要があります。

T = \sum^{x - 1}_{i = 0} A_i + \sum^{y - 1}_{i = x} B_i + \sum^{N - 1}_{i = y} C_i

解法はいくつか考えられますが、本記事では「x を固定して最適な y を探索する」というアプローチを採りました。
この「最適な y を探索する」過程を後方累積最大値で高速化します。

解法の概要

まず A, B, C の前方累積和 P_A, P_B, P_C を求めると、1 \le x \lt y \lt N に対して以下の式が成り立ちます。

T = (P_A[x] - P_A[0]) + (P_B[y] - P_B[x]) + (P_C[N] - P_C[y])

ここで x に依存する項と y に依存する項を分離するために以下のように項を並べ替えます。

T = (P_A[x] - P_B[x]) + (P_B[y] - P_C[y]) + (P_C[N] - P_A[0])

以降は x を固定し、x \lt y \lt N の範囲で P_B[y] - P_C[y] の最大値を求めたいと考えます。
そこで以下のリスト K および後方累積最大値 S_K を定義します。

\begin{align*} K[i] = & P_B[i] - P_C[i] \quad (0 \le i \lt N) \\ S_K[j] = & \max(K[j],\ K[j + 1],\ \dots,\ K[N - 1]) \quad (0 \le j \lt N) \\ S_K[N] = & -\infty \quad (j = N) \end{align*}

制約から y \lt N なので、P_B[N] - P_C[N] は累積最大値の範囲に含めないように注意してください。

以上より、x1 \le x \lt N - 1 の範囲で走査し、x より後方の最大値 S_K[x + 1] を参照することで \mathcal{O}(N) で問題が解けました。

おわりに

後方累積最大値に限らず、添字を扱う際は 0-indexed または 1-indexed のどちらかに統一することをおすすめします。基本的には 0-indexed に揃えるのが楽でしょう。

累積和の場合は初期値 0 を、累積最大値の場合は初期値 -\infty を挿入することで、前方累積・後方累積のどちらでも添字の扱いが一致し、半開区間で考えやすくなります。
ミスを減らす一番の方法は典型に落とし込むことです。この記事では特に添字が複雑になりがちな後方累積において明確な型を示しました。

本記事では深く触れませんでしたが、後方累積和もまた強力なデータ構造です。
左から走査して解けそうな問題で後方の最適値を求める必要が出てきた際は、後方累積和・後方累積最大値が使えるかもしれないと考えてみてください。

Discussion