tiktok_BIT.py

class BIT:
    def __init__(self, n):
        self.n = n
        self.t = [0] * (n + 1)

    def update(self, idx, val):
        idx += 1
        while idx <= self.n:
            self.t[idx] += val
            idx += idx & -idx

    def query(self, idx):
        ans = 0
        idx += 1
        while idx > 0:
            ans += self.t[idx]
            idx -= idx & -idx
        return ans


def fancyPair(pairs, k1, k2):
    n = 200005
    t = BIT(n)
    ans = 0
    pairs.sort()
    lo, hi = 0, -1
    t.update(pairs[0][1], 1)
    for i in range(1, len(pairs)):
        if pairs[0][0] + pairs[i][0] <= k1:
            t.update(pairs[i][1], 1)
            hi = i
        else:
            break
    if hi == -1:
        return 0
    for lo in range(len(pairs)):
        t.update(pairs[lo][1], -1)
        while pairs[hi][0] + pairs[lo][0] > k1:
            t.update(pairs[hi][1], -1)
            hi -= 1
            if hi <= lo:
                return ans
        if k2 >= pairs[lo][1]:
            ans += t.query(k2 - pairs[lo][1])
    return ans

Who will use this problem as OA？