🙆‍♀️

ABC 369 c 等差数列 ランレングス圧縮

2025/01/14に公開

問題概要

問題

長さ N の数列 A = (A_1, A_2, \ldots, A_N)

以下の条件を満たす [l, r] の個数を出力

  • 部分列 (A_l, A_{l+1}, \ldots, A_r) が等差数列

制約

  • 1 \leq N \leq 2 \times 10^5
  • 1 \leq A_i \leq 10^9

入力例

4
3 6 9 3

出力例

8

問題の考え方

等差数列の条件

等差数列の性質を利用して、差分の比較に条件を落とし込む

数列 (A_l, A_{l+1}, \ldots, A_r) が等差数列である条件は、差分が一定であることに同値

これを差分配列 B を用いて表すと:

  • B = (B_1, B_2, \ldots, B_{N-1}), B_i = A_{i+1} - A_i
  • [l, r] が条件を満たすのは、B_l = B_{l+1} = \ldots = B_{r-1} となる場合です。

ランレングス圧縮

数列において、連続する文字と長さのペアとして保存し、同じ値が連続する区間を効率的にまとめる方法
この問題では、差分配列 B に対してランレングス圧縮を適用します。

ランレングス圧縮の手順

  1. 数列を左から順に走査し、現在の値が連続している長さをカウントします。
  2. 値が変化したタイミングで、(値, 長さ) のペアを記録します。
  3. 繰り返します。

例:

  • 入力:B = [2, 2, 2, 3, 3, 1]
  • 圧縮結果:[(2, 3), (3, 2), (1, 1)]

ランレングス圧縮を行うことで、連続する値を1つの区間としてまとめ、効率的に処理できます。


部分区間の数え方

ランレングス圧縮された各区間 (C_i, D_i) において:

  • 長さ D_i の部分区間の数は\frac{D \times (D + 1)}{2}

【補足】なぜ\frac{D \times (D + 1)}{2}になるのか?

長さ D の区間に含まれる部分区間の数を考えます。

  1. 部分区間の長さごとに分けると:

    • 長さ 1 の部分区間:D
    • 長さ 2 の部分区間:D-1
    • 長さ 3 の部分区間:D-2
    • ...
    • 長さ D の部分区間:1
  2. 総和を計算すると:

    • 1 + 2 + 3 + \ldots + D
  3. これは連続する整数の和の公式で計算できます:

    • \text{総和} = \frac{D \times (D + 1)}{2}

実装

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

long long f(long long n){
  return n * (n+1) / 2;
}

int main() {
  int n;
  cin >> n;
  vector<int> a(n);
  for(int &i:a) cin >>i;
  
  if(n==1){
    cout<<1<<endl;
    return 0;
  }
  
  //差分配列に変換
  vector<int> b(n-1);
  for(int i=0;i<n-1;i++){
    b[i]=a[i+1]-a[i];
  }

  //ランレングス圧縮(Run Length Encoding)
  vector<pair<int,int>> rle;
  int cnt=1;
  for(int i=0;i<n-1;++i){
    if(b[i]==b[i+1]){
      cnt++;
    }else{
      rle.emplace_back(b[i],cnt);
      cnt=1;
    }
  }
  rle.emplace_back(b.back(),cnt);

  //答えをランレングス圧縮後の部分区間から算出
  long long ans = n;
  for(auto &[val,len] : rle){
    ans += f(len);
  }

  cout <<ans<<endl;
  return  0;
}

学習内容

  1. 等差数列の処理

    • 等差数列の条件を「差分が一定」という形に変換
  2. ランレングス圧縮

    • 同じ差分を持つ連続区間を圧縮
  3. long long
    int の範囲((2^{31} - 1 = 2,147,483,647))を超えるため、long long が必要

Discussion