find the k-th number of 2 sorted arrays

glassdoor_kthnum2.py

#!/usr/bin/python

def get_less_cnt(buf, pivot):
    left, right = 0, len(buf) - 1
    ret = 0
    while left <= right:
        mid = (left + right) / 2
        if buf[mid] <= pivot:
            ret = max(ret, mid)
            left = mid + 1
        else:
            right = mid - 1
    return ret + 1

def get_kth(buf1, buf2, k):
    left = min(buf1[0], buf2[0])
    right = max(buf1[-1], buf2[-1])
    ret = right
    while left <= right:
        mid = (left + right) / 2
        cnt = get_less_cnt(buf1, mid) + get_less_cnt(buf2, mid)
        if cnt >= k:
            ret = min(ret, mid)
            right = mid - 1
        else:
            left = mid + 1
    return ret

def main():
    buf1 = [1, 3, 5, 7, 9]
    buf2 = [1, 4, 6, 7, 8]
    for i in xrange(1, 11):
        print get_kth(buf1, buf2, i)

if __name__ == '__main__':
    main()