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