公式

E - ちょうどよい温度差 / Just the Right Temperature Difference 解説 by kyopro_friends


\(f(l,r)=\max(A_l,\dots,A_r)-\min(A_l,\dots,A_r)\) と定めます。

考察

考えている区間が広くなると、最大値は大きく、最小値は小さくなるので、 \(f(l,r)\) は大きくなります。つまりある種の単調性があるので、元の問題の「 \(f(l,r)=K\) となる \((l,r)\) は?」の代わりに、「 \(f(l,r)\leq K\) となる \((l,r)\) は?」であれば、尺取法や二分探索を用いて高速に求められるかもしれないと見当をつけることができます。
実際にそれでうまく元の問題が解けることを確かめます。

解法

\(f(l,r) \leq x\) となる \((l,r)\) の個数」を \(g(x)\) とします。元の問題の答えは \(g(K)-g(K-1)\) なので、\(g(x)\) を高速に求めることができれば、元の問題にも答えることができます。

\(A_l,\ldots,A_r\) を ordered multiset で持つことで、\(g(x)\) は尺取法により求めることができ、 \(O(N\log N)\) で問題を解くことができます。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n, k;
  cin >> n >> k;
  vector<int>a(n);
  for(int i=0; i<n; i++) cin >> a[i];

  auto solve = [&](int x){
    if(x < 0){
      return 0LL;
    }
    multiset<int>s;
    int r = 0;
    long long ans = 0;
    for(int l=0; l<n; l++){
      while(r < n){
        int M = a[r], m = a[r];
        if(s.size() > 0){
          M = max(*prev(s.end()), M);
          m = min(*s.begin(), m);
        }
        if(M-m > x)break;
        s.insert(a[r]);
        r++;
      }
      ans += r-l;
      s.erase(s.find(a[l]));
    }
    return ans;
  };

  cout << solve(k) - solve(k-1) << endl;
}

なお、尺取法は区間の操作が「左側を縮める・右側を伸ばす」だけであることから、ordered multiset の代わりに適切に deque を用いることでも最大値・最小値を管理することができ、この方針では \(O(N)\) で問題を解くことができます。

実装例 (Python)

from collections import deque
N, K = map(int, input().split())
A = list(map(int, input().split()))

def solve(x):
  if x < 0:
    return 0
  Mq = deque()
  mq = deque()
  r = 0
  ans = 0
  for l in range(N):
    while r < N:
      M = A[r]
      if len(Mq) > 0:
        M = max(Mq[0][0], M)
      m = A[r]
      if len(mq) > 0:
        m = min(mq[0][0], m)
      if M-m > x:
        break
      while len(Mq) > 0 and Mq[-1][0] < A[r]:
        Mq.pop()
      Mq.append((A[r],r))
      while len(mq) > 0 and mq[-1][0]>A[r]:
        mq.pop()
      mq.append((A[r],r))
      r += 1
    ans += r-l
    if Mq[0][1] == l:
      Mq.popleft()
    if mq[0][1] == l:
      mq.popleft()
  return ans

print(solve(K) - solve(K-1))

投稿日時:
最終更新: