I was solving this sorting problem, and I was wondering if I can get any code review and feedback. You can find this problem online.

#  Find Kth Smallest Element in a Range

#

#  Prompt:    Given an unsorted array of whole integers in the range

#             1000 to 9000, find the Kth smallest element in linear time

#             The array can have duplicates.

#

#  Input:     Unsorted array of whole integers in range of 1000 to 9000

#             Kth smallest element you want to find

#

#  Output:    Kth smalest element in the range

#

#  Example:   array = [1984, 1337, 9000, 8304, 5150, 9000, 8304], k=5

#             output = 8304

#

#  Notes:     What are the time and auxilliary space complexity?



# Time Complexity: O(N)

# Auxiliary Space Complexity: O(N)

def kthSmallest(lst, k):

    if k > len(lst):

        return None

    int_counts = [0 for x in range(8001)]

    for i in range(len(lst)):

        int_counts[lst[i] - 1000] += 1

    cumulative = 

    prev = 0

    for i in range(len(int_counts)):

        cumulative.append(int_counts[i] + prev)

        prev += int_counts[i]

    for i in range(len(lst) - 1, -1, -1):

        if cumulative[lst[i] - 1000] - 1 == k - 1:

            return lst[i]

        cumulative[lst[i] - 1000] -= 1

it passes all of the following test

from cStringIO import StringIO

import sys

import random



# code for capturing print output

#

# directions: capture_print function returns a list of all elements that were

#             printed using print with the function that it is given. Note that

#             the function given to capture_print must be fed using lambda.

#             Example cis provided below

class Capturing(list):

    def __enter__(self):

        self._stdout = sys.stdout

        sys.stdout = self._stringio = StringIO()

        return self



    def __exit__(self, *args):

        self.extend(self._stringio.getvalue().splitlines())

        sys.stdout = self._stdout





# custom assert function to handle tests

# input: count {List} - keeps track out how many tests pass and how many total

#        in the form of a two item array i.e., [0, 0]

# input: name {String} - describes the test

# input: test {Function} - performs a set of operations and returns a boolean

#        indicating if test passed

# output: {None}

def expect(count, name, test):

    if (count is None or not isinstance(count, list) or len(count) != 2):

        count = [0, 0]

    else:

        count[1] += 1



    result = 'false'

    error_msg = None

    try:

        if test():

            result = ' true'

            count[0] += 1

    except Exception, err:

        error_msg = str(err)



    print('  ' + (str(count[1]) + ')   ') + result + ' : ' + name)

    if error_msg is not None:

        print('       ' + error_msg + 'n')







print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + 'nn')



print('Kth Smallest Element Tests')

test_count = [0, 0]





def test():

    example = kthSmallest([1984, 1337, 9000, 8304, 5150, 9000, 8304], 5)

    return example == 8304





expect(test_count, 'should return 8304 for sample input', test)





def test():

    example = kthSmallest([1984, 1337, 9000, 8304, 5150, 9000, 8304], 1)

    return example == 1337





expect(test_count, 'should return 1337 for 1st smallest element with sample input array', test)





def test():

    example = kthSmallest([1984, 1337, 9000, 8304, 5150, 9000, 8304], 10)

    return example == None





expect(test_count, 'should error out when asking for kth smallest when k exceeds size of input array', test)





def test():

    example = kthSmallest([8304], 1)

    return example == 8304





expect(test_count, 'should work for single-element array', test)



def test():

    test_case = 

    for i in range(0, 1000000):

        test_case.append(int(random.random() * 8000) + 1000)

    example = kthSmallest(test_case, 185)

    test_case.sort()

    return example == test_case[184]





expect(test_count, 'should work for a large array', test)



print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + 'nn')

There is no need to have a cumulative array. Using append repeatedly is also no such a great idea, because of the cost of memory reallocation.

def kthSmallest(lst, k):

    if k > len(lst):

        return None

    int_counts = [0 for x in range(8001)]

    for i in range(len(lst)):

        int_counts[lst[i] - 1000] += 1

    cumulative = 0

    for i in range(len(int_counts)):

        cumulative += int_counts[i]

        if cumulative >=  k:

            return i + 1000

Your algorithm requires O(N) additional storage and has more than O(N) complexity (because cumulative.append gets more expensive as cumulative gets larger).

Python has a priority queue implementation, called heapq.

We can use this to implement an algorithm of O(N log K) complexity and O(K) additional storage much more simply, but note that we need to store the negative of the number to turn Python's min-heap into the max-heap we need:

import heapq



def kthSmallest(iterable, k):

    smallest = 

    for value in iterable:

        if (len(smallest) < k):

            heapq.heappush(smallest, -value)

        else:

            heapq.heappushpop(smallest, -value)

    if (len(smallest) < k):

        return None

    return -smallest[0]

We don't even need that loop, as Python provides an equivalent function for us (but this returns a max-heap, so we select the last element instead of negating the first):

import heapq



def kthSmallest(iterable, k):

    smallest = heapq.nsmallest(k, iterable)

    if (len(smallest) < k):

        return None

    return smallest[-1]

Unit tests. Great! Too many reviews don't have tests, so you're already ahead of the game.

We can improve them by giving better messages when tests fail. What we really want to know is what result we got, as we'd like to see how it's different to what was expected. We could enhance our expect() function, but I'll show you how to use the Python unittest module:

import unittest



class TestKthSmallest(unittest.TestCase):



    def test_small(self):

        inputs = [1984, 1337, 9000, 8304, 5150, 9000, 8304]

        # sorted: [1337, 1984, 5150, 8304, 8304, 9000, 9000]

        self.assertEqual(kthSmallest(inputs, 1), 1337)

        self.assertEqual(kthSmallest(inputs, 2), 1984)

        self.assertEqual(kthSmallest(inputs, 3), 5150)

        # now check the last element, and the first n that's too big

        self.assertEqual(kthSmallest(inputs, len(inputs)), 9000)

        self.assertEqual(kthSmallest(inputs, len(inputs)+1), None)



if __name__ == '__main__':

    unittest.main()

These tests will be mercifully silent when they succeed, but if you make them fail, you get to see how they fail:

Traceback (most recent call last):

  File "./201241.py", line 34, in test_small

    self.assertEqual(kthSmallest(inputs, 2), 1337)

AssertionError: 1984 != 1337

For the large test, we'll want to seed the random number generator, to ensure that we're performing the same test every time. A test that sometimes fails is much less help than one that always fails!

def test_large(self):

    random.seed(1)          # ensure the test is reproducible

    inputs = [random.uniform(1000, 9000) for _ in range(10000000)]

    result = kthSmallest(inputs, 185)

    inputs.sort()

    self.assertEqual(result, inputs[184])

As an alternative, we could make this faster by starting with a known order and then shuffling it to get the function input (shuffling should be faster than sorting, unless the random() function is really slow):

def test_large(self):

    random.seed(1)          # ensure the test is reproducible

    inputs = list(range(1, 10000000))

    random.shuffle(inputs)

    result = kthSmallest(inputs, 185)

    self.assertEqual(result, 185)

As a bonus, switching to unittest makes the code work in both Python 2 and Python 3 - that's a Good Thing to have.

apart from @W. Chang's answer, a note on string concatenation

if you have to insert several variables in a string, it's better to use string formatting

from

print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + 'nn')

to

print('PASSED: {} / {}nn'.format(test_count[0], test_count[1]))

it also avoids the use of str each time, giving a clearer idea of output

Here is a minor change to handle arrays with all or mix of -ve elements and possibly all 0's

def kthSmallest(lst,k):

    # Base case

    if k > len(lst):

            return None

    # Case for negative elements, zeros

    m = max(lst)

    if m <= 0:

        m = abs(min(lst))

        if m == 0:

            m = len(lst)



    items= [0] *  (2 * m + 1)

    for i in range(len(lst)):

         new_index =lst[i]+ m

         items[new_index] += 1

    count = 0

    for i, item in enumerate(items):

        count += item

        if count >= k:

            print(i)

            return  i - m