import numpy as np
'''
Problem Description
-------------------
You are given 3 arrays A, B, and P of N elements. A and B contains non-negative
integers less than 10000 and P contains numbers between 0 and 1.
Your task is to compute the expected number of inversions in array C, which
is created using the following rule:
C[i] = A[i] with probability P[i] or B[i] with probability (1 - P[i])
An inversion in C is a pair of elements (C[i], C[j]), i < j, for which
C[i] > C[j].
'''
class fenwicktree(object):
def __init__(self, N):
self.e = np.zeros(N + 2, dtype='float64')
def get(self, i):
if i < 0 or i >= len(self.e) - 1: raise ValueError('i is out of range')
i += 1
answer = 0
while i:
answer += self.e[i]
i -= i & -i
return answer
def update(self, i, d):
if i < 0 or i >= len(self.e) - 1: raise ValueError('i is out of range')
i += 1
while i < len(self.e):
self.e[i] += d
i += i & -i
def count_inversions(A):
'''Counts the number of inversions in A using Fenwick tree.'''
tree = fenwicktree(A.max())
answer = 0
for a in A:
answer += tree.get(a - 1) if a > 0 else 0
tree.update(a, 1.0)
return answer
if __name__ == '__main__':
random_state = np.random.RandomState(42)
# You do not need to try larger arrays to find out the difference
# between the following 2 methods.
N = 14
A = random_state.randint(10000, size=N)
B = random_state.randint(10000, size=N)
C = np.empty_like(A)
P = random_state.rand(N)
I = 0
# Method #1: O(N*log(N))
tree = fenwicktree(max(A.max(), B.max()))
answer = 0.0
for i, (a, b, p) in enumerate(zip(reversed(A), reversed(B), reversed(P))):
answer += p * (tree.get(a - 1) if a > 0 else 0.0)
answer += (1 - p) * (tree.get(b - 1) if b > 0 else 0.0)
tree.update(a, p)
tree.update(b, 1 - p)
print 'Method #1:', answer
# Method #2: O(2**N*N*log(N))
answer = 0.0
for _ in range(2**N):
probability = 1.0
for i, p in enumerate(P):
if (I & (1 << i)) == 0:
probability *= p
C[N - i - 1] = A[i]
else:
probability *= (1 - p)
C[N - i - 1] = B[i]
answer += count_inversions(C) * probability
I += 1
print 'Method #2:', answer