Count candidates with engagement and relevance score are greater than K others

"""
A candidate has two sets of score - engagement and relevance score (es and rs respectively).
Find how many candidates who's engagement and relevance score are greater than K other candidates.

i.e for candidate i --> there must be k candidates where es[i] > es[j] and rs[i] > rs[j]

e.g
es = [1, 2, 3, 4, 5]
rs = [1, 2, 3, 4, 5]

K = 1
Output: 4 since all candidates except the 0th one have at least one other candidate where its engagement and relevance score is greater than another candidate

Implementation should be in O(n log n)
"""

def solve(es, rs, K):
    """
    Solve the candidate ranking problem using a Binary Indexed Tree.

    Time Complexity: O(n log n)
    Space Complexity: O(n)
    """
    n = len(es)

    # Coordinate compression
    es_rank = {v: i+1 for i, v in enumerate(sorted(set(es)))}
    rs_rank = {v: i+1 for i, v in enumerate(sorted(set(rs)))}
    candidates = [(es_rank[e], rs_rank[r]) for e, r in zip(es, rs)]

    # Sort candidates: es ascending, rs descending
    candidates.sort(key=lambda x: (x[0], -x[1]))

    # Initialize Binary Indexed Tree
    bit_size = len(rs_rank) + 1
    BIT = [0] * bit_size

    def update(i):
        # BIT property: update all ranges that include index i
        while i < bit_size:
            BIT[i] += 1
            i += i & -i

    def query(i):
        # BIT property: sum all ranges up to index i
        res = 0
        while i > 0:
            res += BIT[i]
            i -= i & -i
        return res

    result = 0
    i = 0
    while i < n:
        es_val = candidates[i][0]
        same_es = []

        # Group candidates with the same es value
        while i < n and candidates[i][0] == es_val:
            same_es.append(candidates[i])
            i += 1

        # Count candidates with lower rs using BIT
        counts = [query(rs - 1) for _, rs in same_es]

        # Update BIT and count qualifying candidates
        for idx, (_, rs) in enumerate(same_es):
            update(rs)
            if counts[idx] >= K:
                result += 1

    return result


es = [1, 2, 3, 4, 5]
rs = [1, 2, 3, 4, 5]

K = 1

print(solve(es,rs,K)) # 4

es = [5, 4, 3, 2, 1]
rs = [1, 2, 3, 4, 5]

K = 1

print(solve(es,rs,K)) # 0

es = [1, 2, 2, 3, 3, 4, 5]
rs = [1, 2, 7, 2, 3, 4, 5]

K = 3
print(solve(es,rs,3)) # 2

es = [5, 0, 5]
rs = [10, 5, 10]
print(solve(es, rs, 1)) # 2

es = [1, 2, 3, 4, 5]
rs = [5, 4, 3, 2, 1]
print(solve(es, rs, 1)) # 0