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