AGC005-B Minimum Sum 解説

2020-08-23 AtCoder Python 競技プログラミング よりよい計算量シリーズ

\(O(N )\) で解きます。

問題

AtCoder Grand Contest 005 | B - Minimum Sum

概要

$$ \sum _ {l = 1} ^ N \sum _ {r = l} ^ N \mathrm{min} \left( a _ l, a _ {l + 1} , \ldots , a _ r \right) $$

を求めなさい。

解法

editorial 解は平衡二分木を用いたものです。Pythonの標準ライブラリには平衡二分木がないので、別の方法を使う必要があります。（宣伝）

ほかにもいくつか解法があるので、その中から editorial よりも計算量の小さい方法を二つ紹介します。

二つの配列

  • \(\mathit{left}\left[i\right]\) ： \(\mathrm{max} \{\ j\ |\ j < i, a _ j < a _ i \} \)
  • \(\mathit{right}\left[i\right]\) ： \(\mathrm{min} \{\ j\ |\ j > i, a _ j < a _ i \} \)

を用意します。簡単に言うと、それぞれ「自分の左右にある自分より小さいものの中で一番近いもの」のindexを記録したものです。とりあえず、この二つの配列を完成させることを目標にします。

初めに、 \(\mathit{left}\left[i\right] \leftarrow i-1,\ \mathit{right}\left[i\right] \leftarrow i+1 \) で初期化します。 \(a _ i\) の大きい順にこの配列を埋めていきます。 \(a _ i\) を処理しているとき、 \(\mathit{left}\left[i\right]\) より大きく \(\mathit{right}\left[i\right]\) より小さい場所には、 \(a _ {\mathit{left}\left[i\right]}\) よりも大きく、\(a _ {\mathit{right}\left[i\right]}\) よりも大きい要素しか存在していないはずです。

＜ここに図を入れる＞

そこで、 \(a _ {\mathit{left}\left[i\right]} ,\ a _ {\mathit{right}\left[i\right]}\) に対して以下のように更新をしてあげます。

$$ \begin{aligned} \mathit{left}\left[\mathit{right}\left[i\right]\right] &\leftarrow \mathit{right}\left[i\right] \\ \mathit{right}\left[\mathit{left}\left[i\right]\right] &\leftarrow \mathit{left}\left[i\right] \\ \end{aligned} $$

お互いの間にある要素は見る必要がないのでスキップさせるイメージです。これを繰り返していくことで、全 \(a _ i\) に対して正しい \(\mathit{left}\left[i\right],\mathit{right}\left[i\right]\) を求めることができます。

答えに加算していく区間のうち、左端が \(\mathit{left}\left[i\right] + 1,\ \mathit{left}\left[i\right] + 2 ,\ \ldots ,\ i\) のどれかで、右端が \(i,\ i + 1 ,\ \ldots ,\ \mathit{right}\left[i\right] - 1\) のどれかであるようなものは、 区間内の最小値が \(a _ i\) になります。よって最終的な答えは

$$ \sum _ {i = 1} ^ N a _ i ( i - \mathit{left}\left[i\right]) (\mathit{right}\left[i\right] - i) $$

で求めることができます。 一般的にはソートに \(O( N \log N)\) かかってしまいますが、この問題では \(\{a_ i\}\) が順列であることから、バケツソートの要領で \(O(N)\) で並び替えることができます。全体の計算量も \(O(N)\) です。

雑に実装すると両端の要素の更新を行う際におかしくなる可能性があるので、if 文などで適切に処理をしてください。以下の Python のコードでは負のインデックスを与えると末尾からアクセスできる仕様を利用して簡便に実装しています。

ちなみに、更新に Union-Find を利用して \(O(N \alpha (N))\) で解くこともできます。

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

# indexをソート
b = [0] * N
for i in range(N):
    b[a[i] - 1] = i

# 初期化端っこ対策としてひとつ余計に取っている
left = [0] * (N + 1)
right = [0] * (N + 1)
for i in range(N + 1):
    left[i] = i-1
    right[i] = i+1

# 更新
for idx in reversed(b):
    left[right[idx]] = left[idx]
    right[left[idx]] = right[idx]

ans = 0
for i in range(N):
    ans += a[i] * (i - left[i]) * (right[i] - i)

print(ans)

stackの活用

上と同様に \(\mathit{left} ,\ \mathit{right}\) を求めるのですが、今度は先頭の要素から順番に \(\mathit{left}\) を求めていくことにします。

stack \(\mathit{buf}\) を用意し、これに順に \(1,\ 2,\ 3,\ \ldots\) を追加していきます。 \(i\) を追加する際、 \(\mathit{buf}\) 末尾の要素 \(j\) が \(a _ j > a _ i\) である限り pop していきます。 すると、最後に残った末尾 \(j\) は \(a _ j < a _ i\) を満たす最大の \(j\) となっているはずです。（ヒストグラム内最大長方形問題と同じ要領です。 ） \(\mathit{left}\left[i\right] \leftarrow j\) と記録して \(\mathit{buf}\) に \(i\) を push し、次に進みます。

この手順を繰り返すことで、すべての要素について \(\mathit{left}\) を求めることができます。同様の操作を右から行うことで \(\mathit{right}\) も求めることができます。各要素につき高々一度ずつしか push, pop が行われないので、計算量は \(O(N)\) です。

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

left = [0] * (N + 1)
right = [0] * (N + 1)

# 左から
buf = list()
for i in range(N):
    while len(buf) > 0 and a[buf[-1]] > a[i]:
        buf.pop()
    if len(buf) > 0:
        left[i] = buf[-1]
    else:
        left[i] = -1
    buf.append(i)

# 右から
buf = list()
for i in reversed(range(N)):
    while len(buf) > 0 and a[buf[-1]] > a[i]:
        buf.pop()
    if len(buf) > 0:
        right[i] = buf[-1]
    else:
        right[i] = N
    buf.append(i)

ans = 0
for i in range(N):
    ans += a[i] * (i - left[i]) * (right[i] - i)

print(ans)

実装