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

• Mast: Does #!/usr/bin/python even work if it's not on the first line of the file?
• Jason: @Mast It'll not work if you run that on your local machine using ./ to your script name. I guess that's just a prompt (the first line) of HackerRank online editor, and HackerRank will take care of running it correctly behind the scene.

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

• 301_Moved_Permanently: Not sure which of the docstring (and thus your interpretation) or the original code is wrong, but while i > 0: i = i - (i & -i) is NOT the same as range(1, i).
• Oscar Smith: I think it's a docstring problem, what do you want i - (i & -i) to do?
• 301_Moved_Permanently: I don't want it to do anything, it's just not performing i -> 1. For instance, with i = 31 it goes 31 -> 30 -> 28 -> 24 -> 16 -> 0: it removes the least significant bit at each iteration.
• Oscar Smith: oh. Then for the first piece, all I'd change is i=i-... to i-=...