二分探索の練習のためにABC077 C問題を苦戦しながら解く

二分探索の練習のためにABC077のC問題 Snuke Festivalを解きました。
問題を解くポイントと私が苦戦したポイントを解説します。
方針はともかくデバッグに1時間以上かかり、まだまだ二分探索に慣れてないなと思いました。。。
もちろん練習なのでbisectとかは使いません。

atcoder.jp

祭壇のパーツが3種類、A、B、Cとあって、パーツのサイズが $A < B < C$ となるような組み合わせの数を答える問題です。
パーツの個数は最大で $10^5$ なので、単純な総当たりでは無理です。

ポイント1　AとCの両方の制約の影響を受けるパーツBから決める

パーツBを固定すると、パーツAの範囲とパーツCの範囲が決まります。
パーツAやパーツCを固定しても、簡単にはなりません。

ポイント2 二分探索の条件式

パーツBを決めたら、サイズがBより小さいパーツAを選ぶために、パーツAの中でギリギリサイズがBより小さいもの（上限）を探します。
同じく、サイズがBよりも大きいパーツCを選ぶために、パーツCの中でサイズがBより大きいもの（下限）を探します。
二分探索するためにあらかじめパーツAとパーツCのリストはソートしておきます。

二分探索のイメージです。

二分探索のイメージとしては、左端(left)と右端(right)のインデックス値から中央のインデックス値(mid)を計算しながら、leftとrightの範囲を狭めていきます。

ただ、ここからが私がハマったところで、ピッタリとした値を探す二分探索ではなくて、上限や下限を探す二分探索なので、whileの条件やmidの値の出し方に注意しないと、無限ループします。

ここからはデバッガーを駆使して1時間も苦戦してしまいました。。。

これは慣れるしかないのかなぁ。。。1週間後に同じ問題解いたらまた1時間とか苦戦しそうな気がする。。。

N = int(input())
a_arr = list(map(int, input().split()))
a_arr.sort()
b_arr = list(map(int, input().split()))
c_arr = list(map(int, input().split()))
c_arr.sort()


def find_least_idx(num: int, lst: list) -> int:
    n = len(lst)
    left = 0
    right = n - 1
    while left < right:
        mid = (left + right) // 2
        if lst[mid] > num:
            right = mid
            continue
        left = mid + 1
    return right


def find_most_idx(num: int, lst: list) -> int:
    n = len(lst)
    left = 0
    right = n - 1
    while left < right:
        mid = (left + right) // 2 + 1
        if lst[mid] < num:
            left = mid
            continue
        right = mid - 1
    return left


total = 0

for b in b_arr:
    if a_arr[0] >= b:
        continue
    if c_arr[N - 1] <= b:
        continue
    a_most = find_most_idx(b, a_arr)
    c_least = find_least_idx(b, c_arr)
    total += (a_most + 1) * (N - c_least)

print(total)