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):
        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]
        count[1] += 1

    result = 'false'
    error_msg = None
        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)
    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')
  • 3
    \$\begingroup\$ You keep using expect() to test your code. Please look at using pytest or unittest as you might not have a minimum skill level to interview for a coding job. Knowing or having some automated testing experience is a minimum these days. \$\endgroup\$ – C. Harley Aug 9 '18 at 4:52
  • \$\begingroup\$ Are \$\mathcal O(n)\$ for time and space complexities your answer to the note or actual requirements? \$\endgroup\$ – Snowhawk Aug 9 '18 at 7:12
  • \$\begingroup\$ Thanks for including the tests with the code. That definitely made it much easier to experiment with alternative implementations. (Though, as C. Harley observed, you could improve them by using the libraries to help). \$\endgroup\$ – Toby Speight Aug 9 '18 at 16:10

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
| improve this answer | |

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)
            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]
| improve this answer | |

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__':

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)
    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))
    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.

| improve this answer | |

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


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


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

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

| improve this answer | |

Overall good and reviews so far covered key aspects. IMHO, It is always good practice to build safety nets in your code to handle exceptions. One such case would be to handle boundary/corner cases. Declaring in the requirements as below

Unsorted array of whole integers in range of 1000 to 9000 Kth smallest element you want to find

is good, but it should be factored in the program to handle scenarios with input values not in the expected range. By handling I don't mean any specific way. You can explore and choose one as you see fit. But one quick and dirty generalisation of what you wrote that can work with arrays comprising +ve and/or -ve integers and/or 0's might be of interest. Here you go:

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:
            return  i - m    
| improve this answer | |

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.