Official

E - 展示作品の選定 / Selection of Exhibited Works Editorial by admin

Gemini 3.1 Pro (Thinking)

概要

与えられた数列から、隣り合う要素の差の絶対値が \(D\) 以下となるような最長の部分列(要素を順番に選んでできる列)の長さを求める問題です。

考察

まず、素朴な動的計画法(DP)を考えてみましょう。 dp[i] を「\(i\) 番目の作品を最後に選んだときの、選べる作品の最大個数」とします。 すると、以下のような漸化式が成り立ちます。 \(dp[i] = \max(dp[j]) + 1\) (ただし、\(j < i\) かつ \(|H_i - H_j| \leq D\)

しかし、この方法では各 \(i\) について条件を満たす \(j\) をすべて探すため、計算量が \(O(N^2)\) となってしまいます。制約の \(N \leq 2 \times 10^5\) では実行時間制限(TLE)に引っかかってしまいます。

これを高速化するために、DPの持ち方を「インデックス \(i\)」から「作品の評価スコア \(H_i\)」に変えることを考えます。 \(i\) 番目の作品(スコア \(H_i\))を選ぶとき、直前に選ぶ作品のスコア \(H_j\)\(H_i - D \leq H_j \leq H_i + D\) を満たす必要があります。つまり、「これまでに見た作品の中で、スコアが \(H_i - D\) 以上 \(H_i + D\) 以下であるものの DP の最大値」が高速に分かればよいことになります。

これは「区間の最大値取得」と「一点の更新」が高速に行えるデータ構造である セグメント木 (Segment Tree) を使うことで解決できます。 ただし、スコア \(H_i\) は最大 \(10^9\) になるため、そのまま配列のインデックスとして使うとメモリが足りなくなります。そこで、入力に登場するスコアの値だけを抽出してインデックスを振り直す 座標圧縮 というテクニックを併用します。

アルゴリズム

  1. 座標圧縮 入力された評価スコア \(H\) に登場する値を重複を取り除いて昇順に並べた配列 vals を作成します。これにより、最大 \(10^9\) だったスコアを \(0\) から \(N-1\) 以下のインデックスに変換できます。

  2. セグメント木の初期化 区間の最大値を管理するセグメント木を用意します。サイズは vals の要素数以上の \(2\) のべき乗とします。初期値はすべて \(0\) です。

  3. DPの計算 作品を \(1\) 番目から順に見ていきます。現在の作品のスコアを \(h\) とします。

    • 探すべきスコアの範囲は \([h - D, h + D]\) です。
    • 二分探索(bisect_left, bisect_right)を用いて、この範囲が vals のどのインデックス区間 \([l, r]\) に対応するかを求めます。
    • セグメント木を使って、区間 \([l, r]\) の最大値を取得します。この最大値に \(1\) を足した値が、現在の作品を選んだ場合の最大個数 dp となります。
    • 答えとなる最大値を dp で更新します。
    • セグメント木の \(h\) に対応する位置(bisect_left(vals, h))の値を、現在の値と dp のうち大きい方で更新します。
  4. 答えの出力 すべての作品を見終わった後、記録しておいた最大値を出力します。

計算量

  • 時間計算量: \(O(N \log N)\)
    • 座標圧縮のためのソートに \(O(N \log N)\) かかります。
    • 各作品について、二分探索とセグメント木のクエリ・更新にそれぞれ \(O(\log N)\) かかり、これを \(N\) 回行うため \(O(N \log N)\) です。
    • 全体として \(O(N \log N)\) となり、十分に高速です。
  • 空間計算量: \(O(N)\)
    • 座標圧縮用の配列とセグメント木の配列を保持するため、作品の個数 \(N\) に比例したメモリを使用します。

実装のポイント

  • 非再帰のセグメント木: Pythonでは関数呼び出しのオーバーヘッドが大きいため、再帰関数を用いたセグメント木の実装は遅くなることがあります。解答コードのように 1次元配列と while ループを用いたボトムアップ(非再帰)の実装にすると、より高速に動作します。

  • 区間の存在判定: 二分探索で求めた区間 \([l, r]\) において、条件を満たすスコアが一つも存在しない場合(\(l > r\) となる場合)があります。このときはセグメント木にクエリを投げず、最大値を \(0\) として扱う処理を入れる必要があります。

    ソースコード

import sys
from bisect import bisect_left, bisect_right

def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    D = int(data[1])
    H = [int(x) for x in data[2:]]
    
    vals = sorted(list(set(H)))
    M = len(vals)
    
    size = 1
    while size < M:
        size *= 2
    
    tree = [0] * (2 * size)
    
    def query(l, r):
        l += size
        r += size
        res = 0
        while l <= r:
            if l % 2 == 1:
                if tree[l] > res:
                    res = tree[l]
                l += 1
            if r % 2 == 0:
                if tree[r] > res:
                    res = tree[r]
                r -= 1
            l //= 2
            r //= 2
        return res

    def update(pos, value):
        pos += size
        if value > tree[pos]:
            tree[pos] = value
            pos //= 2
            while pos > 0:
                left_child = tree[2 * pos]
                right_child = tree[2 * pos + 1]
                tree[pos] = left_child if left_child > right_child else right_child
                pos //= 2

    ans = 0
    for h in H:
        l = bisect_left(vals, h - D)
        r = bisect_right(vals, h + D) - 1
        
        if l <= r:
            max_val = query(l, r)
        else:
            max_val = 0
            
        dp = max_val + 1
        if dp > ans:
            ans = dp
            
        pos = bisect_left(vals, h)
        update(pos, dp)
        
    print(ans)

if __name__ == '__main__':
    solve()

この解説は gemini-3.1-pro-thinking によって生成されました。

posted:
last update: