Find Kth Smallest Element in Union of Sorted Lists

find_kth_smallest.py

# Copyright (c) 2014 Rahul

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import unittest
import random
import functools


def find_kth_smallest(list_a, list_b, k):
    low_a = 0
    high_a = len(list_a)
    low_b = 0
    high_b = len(list_b)

    while high_a - low_a and high_b - low_b:
        mid_a = (high_a + low_a) / 2
        mid_b = (high_b + low_b) / 2

        elems_a = (mid_a - low_a) + 1
        elems_b = (mid_b - low_b) + 1
        elems = elems_a + elems_b
        if list_a[mid_a] > list_b[mid_b]:
            if k < elems:
                high_a = mid_a
            else:
                low_b = mid_b + 1
                k -= elems_b
        else:
            if k < elems:
                high_b = mid_b
            else:
                low_a = mid_a + 1
                k -= elems_a

    if high_a - low_a:
        return list_a[low_a + k - 1]
    else:
        return list_b[low_b + k - 1]


class FindKthSmallestTest(unittest.TestCase):

    def get_random_sorted_list(self, num):
        return sorted((random.randint(0, 10000) for _ in xrange(num)))

    def run_test_multiple_times(f):
        @functools.wraps(f)
        def func(*args, **kwargs):
            for _ in xrange(10):
                f(*args, **kwargs)
        return func

    @run_test_multiple_times
    def test_find_kth_smallest(self):
        list_a = self.get_random_sorted_list(100000)
        list_b = self.get_random_sorted_list(100000)
        list_c = list_a + list_b
        list_c.sort()

        k = random.randint(1, len(list_c) - 1)

        expected_result = list_c[k - 1]
        returned_result = find_kth_smallest(list_a, list_b, k)

        self.assertEqual(expected_result, returned_result)


if __name__ == '__main__':
    unittest.main()