findMedian.py

class Solution:
    # @param A : list of list of integers
    # @return an integer

    def edge(self, A, x, first=None):
        left = first or 0
        right = len(A)-1
        found = False

        if x < A[0]:
            return 0, found

        while (left <= right):
            m = (left + right)/2
            e = A[m]
            if e > x:
                right = m - 1
            elif e < x:
                left = m + 1
            else:
                found = True
                if first is None:
                    right = m - 1
                else:
                    left = m + 1
        return left, found

    def edges(self, A, x):
        l, r = 0, 0
        for a in A:
            ll, f = self.edge(a, x)
            l = l + ll
            if f:
                rr, _ = self.edge(a, x, ll)
                r = r + rr
            else:
                r = r + ll
        return l, r

    def findMedian(self, A):
        l = len(A)
        L = len(A[0])
        i = L/2
        k = L*l / 2

        left = min((a[i] for a in A))
        right = max((a[i] for a in A))

        if left == right:
            return left

        while left <= right:
            m = (left+right)/2
            lK, rK = self.edges(A, m)

            if lK <= k < rK:
                return m

            if rK > k:
                right = m - 1
            elif lK <= k:
                left = m + 1

print Solution().findMedian([[1, 3, 5],[2, 6, 9],[3, 6, 9]]) == 5 # [1, 2, 3, 3, 5, 6, 6, 9, 9] A[4] = 5
print Solution().findMedian([[1,6,6], [6,6,6], [6,6,9]]) == 6
print Solution().findMedian([[1], [3], [3]]) == 3
print Solution().findMedian([[3], [5], [1], [2], [2], [2], [5]]) == 2 # [1, 2, 2, 2, 3, 5, 5] A[3] = 2
print Solution().findMedian([[9, 10, 10, 13, 14, 15, 16, 16, 16, 17, 18], [1, 4, 9, 14, 16, 18, 19, 22, 26, 26, 27],  [4, 6, 7, 10, 14, 20, 21, 23, 24, 27, 28]]) == 16