えびちゃんの日記

えびちゃん(競プロ)の日記です。

Python でも set::lower_bound っぽいことがしたい

したくない人は読まなくていいです。

TL; DR

vEB

やりたいこと

Pythonset は「x 以上の最小値は?」「x より大きい最小値は?」といったクエリに答えられないので困ります。それをしたいです。x は整数としちゃいます。

特に競プロの文脈の話なので、外部ライブラリにこれこれがあるよーというのは受け付けません。

用語について

背景 1

長いのでたたんじゃいます

lower_boundC++ の用語です。upper_bound もあります。 この言い回しをする気持ちに慣れていない人もいるかもなので、念のため説明しておきます。

以下のような配列を考えます(一旦 set ではないことにして、重複を許します)。

a = {0, 2, 3, 4, 4, 4, 5}

このとき、x と等しい*1区間の下限と上限の位置が、それぞれ lower_bound(a.begin(), a.end(), x)upper_bound(a.begin(), a.end(), x) です。 両方をまとめて返す関数 equal_range(a.begin(), a.end(), x) もあります。

図示すると、次のような感じになります。

a = {0, 2, 3, 6, 6, 6, 8}
     [                     <= lower_bound(0)
        )                  <= upper_bound(0)
              [            <= lower_bound(4)
              )            <= upper_bound(4)
              [            <= lower_bound(6)
                        )  <= upper_bound(6)

upper_bound の方を含まない半開区間なので、それを意識した表記にしています。

名前に関する気持ちはこんな感じなのですが、実際には、「x 以上の最小値が *lower_bound(x)」「x より大きい最小値が *upper_bound(x)」として使われがちです。返り値は値ではなく位置であり、a.end() が返ってくることも当然あるので、何も確認せずに * をするのをおまじないとして覚えているとよくないですね*2

返ってくるのは位置なので、x 以上 y 未満の値の個数が upper_bound(y) - lower_bound(x) で求められたりするわけですね。

set においても話は同じで、位置が返ってきます。 std::setイテレータは引き算ができないので、(std::distance を使って線形時間かけないと)距離は求められません。

背景 2

まだある

前述の通り、bound というのは値ではなく位置を返します。 今回は値を知りたい話なので、別の呼び方をしたいです。

x より大きい最小値」は、x の successor と呼ばれることが多い気がします*3。 「x より小さい最大値」は x の predecessor と呼ばれます。

x 以上の最小値」であるところの *lower_bound に相当する呼び方がなくて困ります...。 こちらが successor であれば、「x より大きい最小値」を指して proper successor とか strict successor とか呼べて便利そうな気もします。狭義単調増加とか部分集合とか真に大きいみたいな言い回しですね。

この記事中では「以上」の方を successor と呼ぶよーとしてもよいのですが、混乱しそうなのでそういうことは控えます。

それから、successor と predecessor を総称した言い方があるとうれしいのですが、あまりなさそうです。 両方向のクエリが来る状況設定でも「predecessor query problem」とか呼ばれがちな気がします。 個人的には(bidirectional とかみたいな気持ちで)bicessor とか呼べたらうれしい気もするのですが...

定義

といった事情があるので、次のように呼んじゃいます。

  • min_gt(x)x より大きい (greater than) 最小値
  • min_ge(x)x 以上 (greater than or equal to) 最小値
  • max_lt(x)x より小さい (less than) 最大値
  • max_le(x)x 以下 (less than or equal to) 最大値

それに加えて、次のものも欲しいのでつくります。

  • add(x)x を追加する
  • remove(x)x を削除する
  • has(x)x が含まれていれば True、でなければ False

つくりかた

実装するまでで疲れちゃったので説明はしません。ごめんね*4

Erik Demaine 大先生の講義を聞くといいです。

www.youtube.com

実装例はこちらです。\(u=2\) になったら再帰せずにがちゃがちゃしています。 \(u\) がなにかという説明も疲れちゃったので割愛します。ごめんね。

atcoder.jp

TL 2 sec のところ 1998 ms でした...。 たぶん Python 慣れしている人ならもっと速い実装をできるんだと思います。

(追記)ABC 228 D は TL が長く、たすかります。

atcoder.jp

関連記事

qiita.com

上記の実装例の問題は、この記事にある方法「別の解法を考える」によって簡単に解くことができます。

atcoder.jp

おわり

いかがでしたか?

C++ や Rust などの言語を使うと楽に書ける場合があります。

*1:ここでの「等しい」は equal ではなく equivalent です。詳細はここでは省きます。

*2:適当に大きい値(こういうのを番兵と呼びます)を末尾に入れておくのもありますが、適切でない番兵を入れてバグらせる初心者もよくいますね。

*3:そうした形式のクエリのことを successor query と呼んだりします。

*4:申し訳ないです...。