Count the number of inversions using Binary Indexed Tree in an array

Question

Original Problem: HackerRank

I am trying to count the number of inversions using Binary Indexed Tree in an array. I've tried to optimise my code as much as I can. However, I still got TLE for the last 3 test cases. Any ideas that I can further optimise my code?

# Enter your code here. Read input from STDIN. Print output to STDOUT
#!/usr/bin/python

from copy import deepcopy



class BinaryIndexedTree(object):
    """Binary Indexed Tree
    """
    def __init__(self, n):
        self.n = n + 1
        self.bit = [0] * self.n

    def update(self, i, k):
        """Adds k to element with index i
        """
        while  i <= self.n - 1:
            self.bit[i] += k
            i = i + (i & -i)

    def count(self, i):
        """Returns the sum from index 1 to i
        """
        total = 0
        while i > 0:
            total += self.bit[i]
            i = i - (i & -i)
        return total


def binary_search(arr, target):
    """Binary Search
    """
    left, right = 0, len(arr) - 1
    while left <= right:
        mid = left + ((right - left) >> 1)
        if target == arr[mid]:
            return mid
        elif target < arr[mid]:
            right = mid - 1
        else:
            left = mid + 1

    # If not found, return -1.
    return -1


T = input()
for iterate in xrange(T):
    n = input()
    q = [ int( i ) for i in raw_input().strip().split() ]
    answer = 0
    # Write code to compute answer using x, a and answer

    # Build a Binary Indexed Tree.
    bit = BinaryIndexedTree(n)

    # Copy q and sort it.
    arr = sorted(deepcopy(q))

    # index array.
    index = map(lambda t: binary_search(arr, t) + 1, q)

    # Loop.
    for i in xrange(n - 1, -1, -1):
        answer += bit.count(index[i])
        bit.update(index[i] + 1, 1)
    print answer

Answer 1

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

bit has a well-known interpretation in programming which distracts from the intended interpretation here. IMO even something as generic as data would be better.

---

    def update(self, i, k):
        """Adds k to element with index i
        """
        while  i <= self.n - 1:
            self.bit[i] += k
            i = i + (i & -i)

The loop condition would be more familiar as while i < self.n.

I don't see any reason to avoid +=.

A reference to explain the data structure would be helpful, because otherwise this is quite mysterious.

---

        mid = left + ((right - left) >> 1)

The reason for this particular formulation of the midpoint is to avoid overflow in languages with fixed-width integer types. Since that's not an issue in Python, I'd favour the more straightforward mid = (left + right) >> 1. It might even be a tiny bit faster.

---

    q = [ int( i ) for i in raw_input().strip().split() ]

    index = map(lambda t: binary_search(arr, t) + 1, q)

Why the inconsistency? I think it would be more Pythonic to use comprehensions for both.

---

    # Copy q and sort it.
    arr = sorted(deepcopy(q))

    # index array.
    index = map(lambda t: binary_search(arr, t) + 1, q)

This seems like a rather heavyweight approach. Why not just

    index = sorted((val, idx) for idx, val in enumerate(q))

?

---

Although asymptotically I don't see any reason why this would be \$\omega(n \lg n)\$, it has a few \$\Theta(n \lg n)\$ stages. The standard algorithm for this problem, which is essentially merge sort, has the same asymptotic complexity but probably hides a much smaller constant behind the \$\Theta\$ / \$O\$.

Answer 2

I'm pretty sure you can replace

def count(self, i):
    """Returns the sum from index 1 to i
    """
    total = 0
    while i > 0:
        total += self.bit[i]
        i = i - (i & -i)
    return total

with

def count(self, i):
    return sum(self.bit[i] for i in range(1,i))

also,

mid = left + ((right - left) >> 1)

is the same as

mid = (right + left) >> 1