Codesprint Direct Connections n^1.5 Solution

import math
cases = int(raw_input())
def binsearch(li,score):
  left,right = 0,len(li)-1
  if right < left: return 0
  while right-left > 1:
    ave = int((left+right)/2)
    if score < li[ave][1]: right = ave # Ascending
    else: left = ave
  if score <= li[left][1]: return left
  elif score <= li[right][1]: return right
  else: return right + 1
for case in range(cases):
  citiesnum = int(raw_input())
  coords = [int(x) for x in raw_input().split(' ')[:citiesnum]]
  pops = [int(x) for x in raw_input().split(' ')[:citiesnum]]
  # Bypop: cities sorted by population (high to low)
  bypop = []
  for x in range(citiesnum): bypop.append([pops[x],coords[x]]) # Pop = [i][0], coord = [i][1]
  bypop.sort(reverse=True,key=lambda i: i[0])
  # Bycoord: cities sorted by coordinates (left to right)
  bycoord = sorted(bypop,key=lambda i: i[1])
  batchsize = int(math.sqrt(citiesnum))
  # Cities sorted by population in batches (ie. [a b c d e f g h i j] -> [[a b c][d e f][g h i][j]]; this is done to reduce
  # the time of some operations from O(n) to O(sqrt(n)), and is the reason why this algorithm outperforms the naive n^2 one so greatly
  batches = []
  # Running total of batch coordinates (eg. 1 2 3 4 5 6 7 8 9 10 -> 6 21 45 55
  macrosums = []
  # Running total of city coordinates within each batch (eg. 1 2 3 4 5 6 7 8 9 10 -> [1 3 6][4 9 15][7 15 24][10]
  microsums = []
  for i in range(0,citiesnum,batchsize):
     microsums.append([])
     batches.append([])
     newsum = 0
     for j in range(i,min(i+batchsize,citiesnum)):
       batches[-1].append(bycoord[j])
       newsum += bycoord[j][1]
       microsums[-1].append(newsum)
     macrosums.append((0 if i == 0 else macrosums[-1]) + newsum)
  total = 0
  remaining = citiesnum
  for i in range(citiesnum-1):
    # Find ith largest city in bycoord, then in batches (outer,inner), get running coordinate sum
    bycoord_place = binsearch(bycoord,bypop[i][1])
    outer = int(bycoord_place / batchsize)
    inner = binsearch(batches[outer],bypop[i][1])
    sumat = (0 if outer == 0 else macrosums[outer - 1]) + microsums[outer][inner]
    # Find previous city in coordinate order in batches
    if outer > 0 or inner > 0:
      innerprev,outerprev = inner - 1,outer
      while innerprev < 0 and outerprev >= 0:
        outerprev -= 1
        innerprev = len(batches[outerprev])-1
      if outerprev < 0: sumprev = 0
      else: sumprev = (0 if outerprev == 0 else macrosums[outerprev - 1]) + microsums[outerprev][innerprev]
    else: outerprev,innerprev,sumprev = -1,-1,0
    place = inner
    for j in range(outer): place += len(batches[j])
    # The basic algorithm works like this:
    # Take some cities:          1 2 3 5  7  9  12 15
    # Get their coordinate sums: 1 3 6 11 18 27 39 54
    # We're looking to find the sum of all distances from, say, the city at 9 to all the others.
    # The answer is (15-9)+(12-9)+(9-7)+(9-5)+(9-3)+(9-2)+(9-1)
    # Rewriting that, we get (15+12) - (1+2+3+5+7) + (9+9+9+9+9) - (9+9)
    # Or (let S(n) be the sum of all cities left of n): S(15) - S(9) - S(7) + 9(10-7), which is what the nt = line below calculates
    # The 10-7 arises from the fact that there are 5 cities to the left, which are added, and 2 to the right, which are subtracted.
    # An alternative way of seeing that is to subtract all 7 other cities (-7) and then add all those to the left twice (+10)
    nt = macrosums[-1] - sumat - (0 if bycoord_place == 0 else sumprev) + bycoord[bycoord_place][1] * (place * 2 + 1 - remaining)
    # Remove the used city from the batches because we don't need to count it again and adjust the sums to account for this
    # eg. ... 18 27 39 54 -> ... 18 30 45
    # It's the adjusting of the sums that makes buckets necessary - if it weren't for buckets, adjusting sums would be an O(n) process,
    # making the entire algorithm O(n^2); buckets reduce the runtime to O(sqrt(n))
    for j in range(inner,len(batches[outer])): microsums[outer][j] -= bycoord[bycoord_place][1]
    for j in range(outer,len(batches)): macrosums[j] -= bycoord[bycoord_place][1]
    batches[outer].pop(inner)
    microsums[outer].pop(inner)
    remaining -= 1
    total += nt * bypop[i][0]
    # Mandated by problem specifications
    total = total % 1000000007
  print total