Official

E - Max Min Editorial by Nyaan


まず、愚直に \(L,R\) を全探索すれば \(\mathrm{O}(N^3)\) あるいは \(\mathrm{O}(N^2)\) でこの問題を解くことができますが、 \(N \leq 2\times 10^5\) なので TLE してしまいます。そこで、アルゴリズムを上手く用いて \(\mathrm{O}(N)\) でこの問題を解く方法を紹介します。

まず、この問題は条件を満たす整数の組 \((L,R)\) の数え上げですが、少し見方を変えると次のように言い換えられます。

\(A\) の連続部分列 \(A_{L}, A_{L+1}, \dots, A_{R}\) のうち、最大値が \(X\), 最小値が \(Y\) となるものは何個あるか。

このように \((L,R)\) の数え上げを区間の数え上げとみなすと見通しが良くなります。
ただ、これでもまだ以下に述べる通り厄介な要素が残っているので、条件を満たす連続部分列がどのようなものかを考察して問題を単純化できないか考えてみましょう。

\[A = (4,2,5,4,3,4,2), X=4, Y=2\]

の場合を考えましょう。

この例では \(A_3 = 5\) が厄介です。\(A_3 = 5\)\(X = 4\) よりも大きいので、\(A_3\) を含む連続部分列はすべて条件を満たしません。このような数を残したまま考察を進めると少し面倒な場合分けが必要になります。
少し考察を進めると、よくよく考えると \(A_3\) を含む連続部分列が条件を満たさないということは、列から \(A_3\) を取り除いて前後で分割して、それぞれの列に対して条件を満たす列を数え上げてもよい、ということになります。すなわち、\(A\) の前 \(2\) つを \(A_1\), 後ろ \(4\) つを \(A_2\) として

\[ A_1 = (4, 2), X=4, Y=2\]

\[A_2 = (4,3,4,2), X=4, Y=2\]

\(2\) つの場合について問題を解いて答えを足し合わせたものが全体の答えとなります。

一般の場合についても同様です。 \(A_i\)\(X \leq A_i \leq Y\) を満たさない場合は、\(A_i\) を取り除いて列を分割して、分割された列に対して数え上げを行えばよいです。

さて、言い換えの結果、分割された列を \(1\) つ取ってきて \(B\) としたとき、次の部分問題 1 が高速に解ければよいことになります。

部分問題 1

\(B\) の連続部分列 \(B_{L}, B_{L+1}, \dots, B_{R}\) のうち、最大値が \(X\), 最小値が \(Y\) となるものは何個あるか。ただし \(B\) の要素は \(Y\) 以上 \(X\) 以下である。

元の問題と比べて太線で強調した条件が追加されています。実はこの条件をうまく利用することで、解くべき部分問題は更に言い換えられます。
列の要素は \(Y\) 以上 \(X\) 以下という条件がすでに満たされているので、「最大値が \(X\), 最小値が \(Y\) となる」という条件は「\(X,Y\) がともに列に含まれる」という条件と等価になっています。
よって次の部分問題 2 を高速に計算できればよいです。

部分問題 2

\(B\) の連続部分列 \(B_{L}, B_{L+1}, \dots, B_{R}\) のうち、\(X, Y\) がともに含まれる列 は何個あるか。

上の問題を線形時間、すなわち \(\mathrm{O}(|B|)\) で解ければ、\(A\) を分割した列の長さの 和は \(N\) を超えないので、全体の計算量は \(\mathrm{O}(\sum_{B \in (A を分割した列の集合)} B) = \mathrm{O}(N)\) となります。

部分問題 2 を高速に解く方法はいくつかあります。1 つずつ説明していきましょう。

解法 1 : 尺取り法

以下では \(B_L,B_{L+1},\dots,B_R\) を 「区間 \([L,R]\) 」と表現します。

部分問題 2 では、次のような性質が成り立ちます。

区間 \([k,l]\) が区間 \([i,j]\) を含んでいる、すなわち \(i,j,k,l\)\(1 \leq k \leq i \leq j \leq l \leq |B|\) の関係にあるとき、区間 \([i,j]\) が条件を満たすならば \([k,l]\) が条件を満たす。

このような場合は次のようなアルゴリズムによって、条件を満たす区間を数え上げることができます。

  • はじめ、\(L=R=1\) とする。
  • \(L\)\(N\) 以下の間、次の操作を行う。
    • 区間 \([L,R]\) が条件を満たさない間、\(R\)\(1\) を足す。
    • 条件を満たす \([L,R]\) が見つかれば、区間 \([L,x]\) \((R\leq x \leq B)\) はすべて条件を満たすことになるので、その分を答えに加算する。
    • \(L\)\(1\) を加える。

このように両端を \(1\) ずつズラしてすべての極小な区間を走査するアルゴリズムを 尺取り法 と呼びます。\(L,R\) をインクリメントする回数が \(|B|\) 回で抑えられているので、条件判定を \(\mathrm{O}(1)\) で処理できれば全体で \(\mathrm{O}(|B|)\) となります。

Python による実装例は次の通りです。

N, X, Y = map(int, input().split())
A = list(map(int, input().split()))

def calc(B):
  i, j, countX, countY, res = 0, 0, 0, 0, 0
  while i != len(B):
    while j != len(B) and (countX == 0 or countY == 0):
      if B[j] == X:
        countX += 1
      if B[j] == Y:
        countY += 1
      j += 1
    if countX > 0 and countY > 0:
      res += len(B) + 1 - j
    if B[i] == X:
      countX -= 1
    if B[i] == Y:
      countY -= 1
    i += 1
  return res


i, ans = 0, 0
while i != N:
  B = []
  while i != N and Y <= A[i] <= X:
    B.append(A[i])
    i += 1
  if len(B) != 0:
    ans += calc(B)
  else:
    i += 1
print(ans)

解法 2 : 包除原理

2 番目の想定解として 包除原理 を紹介します。
包除原理とは数え上げで非常によく出てくるアルゴリズムです。特に条件が \(2\) 個の場合は次の \(2\) 種類の式が知られています。

\[ \begin{aligned} &(P か Q の少なくとも一方が真である集合) \\ &=(P が真である集合) + (Q が真である集合)-(P,Q 両方が真である集合) \end{aligned} \]

\[ \begin{aligned} &(P か Q の両方が真である集合) \\ &=(全集合) - (P が偽である集合) - (Q が偽である集合) + (P,Q 両方が偽である集合) \end{aligned} \]

高校数学の「集合と論理」で学習した内容を覚えている方は、ベン図を書いてみるとわかりやすいと思います。

求めたいのは「 \(X\) を含み、かつ \(Y\) を含む列」の個数ですが、この条件に包除原理を適用すると

\[ \begin{aligned} &(全ての列)\\ &-(X を含まない列)\\ &-(Y を含まない列)\\ &+(X,Y ともに含まない列) \end{aligned} \]

という式で表せます。「~~を含まない列」は実は「~~を含む列」よりも数えやすいので、このように包除原理を適用することでこの問題を解くことができます。(詳しいアルゴリズムは実装例を参照ください。)計算量は \(\mathrm{O}(|B|)\) です。

def f(a):
  ret, s = 0, 0
  for x in a:
    if x == 1:
      ret += s * (s + 1) // 2
      s = 0
    else:
      s += 1
  ret += s * (s + 1) // 2
  return ret


N, X, Y = map(int, input().split())
A = list(map(int, input().split()))
__, X_, _Y, XY = [0] * N, [0] * N, [0] * N, [0] * N

for i in range(N):
  if not (Y <= A[i] and A[i] <= X):
    __[i], X_[i], _Y[i], XY[i] = 1, 1, 1, 1
  if A[i] == X:
    X_[i], XY[i] = 1, 1
  if A[i] == Y:
    _Y[i], XY[i] = 1, 1

print(f(__) - f(X_) - f(_Y) + f(XY))

他にも次のような \(\mathrm{O}(N \log N)\) 解法も考えられて、計算量は少し劣りますがこれでも十分高速なので AC できます。

  • すべての \(i\) に対して「\(i\) より右で \(X,Y\) が初めて出現する箇所は?」というクエリを処理できれば良い。これは \(X,Y\) の出現箇所を累積和や Segment Tree で記録して二分探索を行ったり、 std::map で管理して lower_bound を使えば クエリ当たり \(\mathrm{O}(\log N)\) で処理できる。

posted:
last update: