K’th SmallestElement in Unsorted Array in python 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 &gt; 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]) + '\n\n')

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]) + '\n\n')
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 &gt; 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 &gt;=  k:
            return i + 1000