pythonでlower_bound,upper_boundっぽいもの

 例えばソートされたリストAの中のどの位置に要素nがあるか知りたいとき、二分探索という手法を使えば\(O(\log(N))\)で計算できることはよく知られています。
 しかし、二分探索ではnがAのリスト内に複数含まれるときどこのnのindexが返ってくるのかわからないという弱点があります。例えば

A=[1,2,3,4,4,4,4,4,4,5,6,7,8,9]
n=4

 と与えられていると、6個ある4のどのindexを取得するかはシンプルな二分探索だと指定できません。この仕様はしばしば不便です。
 これを改良した函数であるlower_bound(),upper_bound()を作ったので書いておきます(ちなみにlower_bound,upper_boundはC++の同名の組み込み函数に似せてます)。
 ただし以下の函数はAの要素がintであるときだけ正しい動作が保証されることに注意してください

 (bisectをimportすればわざわざ自分で定義しなくても似たようなことできるけど、自分で書いてみたら勉強になるよね! 嘘です。本当は書いてからbisectというパッケージの存在に気づきました)

lower_bound

def lower_bound(A,n):
        #半整数にする。
        n=n-0.5
        left=0
        right=len(A)-1
        lenA=right-left+1
        while lenA>0:
            if A[left+lenA//2]>n:
                right=left+lenA//2-1
                lenA=right-left+1
            elif A[left+lenA//2]<n:
                left=left+lenA//2+1
                lenA=right-left+1
        #半整数なので探索しても見つかりません。このときleftが求めていたものになっています。
        return left

 ソートされたリストA及びnが与えられたとき、函数lower_boundはA内にあるn以上の要素の中で一番左にあるもののindexを返します。

 ここでやっていることはほぼ二分探索ですが、探索する値をnではなくてn-0.5としています。こうすることで上手く求められています。ただし半整数を使って実装してあるので、この函数はAの要素がすべてintでなければ正しい動作が保証されません。
 やっていることとしては、半整数n-0.5で二分探索するとn-0.5はAの中には存在しないので、ループのどこかの段階でleftがn-0.5以下の最大の数、rightがn-0.5以上の最小の数という状態になります。そこから更に1回ループを実行するとleft,rightが反転してループを脱出し、そのときのleft(=right+1)が求めたいものになっています。

upper_bound

def upper_bound(A,n):
        #半整数にする。
        n=n+0.5
        left=0
        right=len(A)-1
        lenA=right-left+1
        olenA=lenA
        while lenA>0:
            if A[left+lenA//2]>n:
                right=left+lenA//2-1
                lenA=right-left+1
            elif A[left+lenA//2]<n:
                left=left+lenA//2+1
                lenA=right-left+1
        return left

  ソートされたリストA及びnが与えられたとき、函数upper_boundはA内にあるnより真に大きい要素の中で一番左にあるもののindexを返します。

 やってることはupper_boundとほぼ同じです。nを半整数にするところだけが違います。

 いずれの函数も、nがAのどの要素よりも小さいときには0、nがAのどの要素よりも大きいときはlen(A)を返します。

計算量

 lower_boundもupper_boundのいずれも二分探索と似た手法で実装しているのでどちらも\(O(\log(N))\)で実行できます。

これで解ける問題

AtCoder Beginner Contest 077 C – Snuke Festival

コメント

タイトルとURLをコピーしました