Select nth smallest distinct integer from inputs

Question

Kata in question

Given a list of integers, return the nth smallest integer in the list. Only distinct elements should be considered when calculating the answer. n will always be positive (n > 0)

If the nth small integer doesn't exist, return -1 (C++) / None (Python) / nil (Ruby) / null (JavaScript).

Notes:

• "indexing" starts from 1
• huge lists (of 1 million elements) will be tested

Examples
nth_smallest([1, 3, 4, 5], 7) ==> None # n is more than the size of the list
nth_smallest([4, 3, 4, 5], 4) ==> None # 4th smallest integer doesn't exist
nth_smallest([45, -10, 4, 5, 4], 4) ==> 45 # 4th smallest integer is 45 If you get a timeout, just try to resubmit your solution. However, if you always get a timeout, review your code.

I wrote several functions to try to solve this.
My first attempt was:

def nS1(arr = z, n = 16):
    st = set(arr)
    return sorted(st)[n-1] if n <= len(st) else None

My second:

def nS2(arr = z, n = 16):
    st = set()
    count = 0
    for i in sorted(arr):
        if i not in st:
            count += 1
            if count == n:
                return i
            st.add(i)
    return None

I even tried implementing quickselect:

#Given a list, it modifies it so that the element at `pvtIdx` (pivot index) is the `pvtIdx` smallest element.
def partition(lst, lft, rght, pvtIdx):  
    pvtVal = lst[pvtIdx]    #The value of the pivot element that would be used in comparison.
    lst[pvtIdx], lst[rght] = lst[rght], lst[pvtIdx]     #Swap `lst[rght]` and `lst[pvtIdx]`.
    strIdx = lft    #The store index that contains the location that is partitioned.
    for i in range(lft, rght):  #Iterate through the list.
        if lst[i] < pvtVal:     #If the current element is less than the pivot element.
            lst[i], lst[strIdx] = lst[strIdx], lst[i]   #Swap the current element and the partitioner.
            strIdx += 1     #Increment the partitioner.
    lst[rght], lst[strIdx] = lst[strIdx], lst[rght]     #Swap the pivot element and the partitioner.
    #The list is now partitioned into elements < the pivot elements and elements > the pivot element around the partition location.
    return strIdx   #Return the partition location.

def select(lst, lft, rght, k):
    if lft == rght:     #Return the sole element of the list if it is already sorted.
        return lst[lft]
    pvtIdx = lft + int(random()*(rght - lft))       #Generate a random pivot index between `lft` and `rght` (both inclusive).
    pvtIdx = partition(lst, lft, rght, pvtIdx)  #The index of the pivot value in it's sorted position.
    if k == pvtIdx:     #If that index corresponds to the desired index.
        return lst[k]
    elif k < pvtIdx:    #Insert another element to its sorted position in the partition of the list that the desired element resides in.
        return select(lst, lft, pvtIdx - 1, k)
    else:
        return select(lst, pvtIdx + 1, rght, k)     #Insert another element to its sorted position in the partition of the list that the desired element resides in.

def nS3(lst = z, k = 16):
    st = set(lst)
    ln = len(st)
    return None if k > ln else select(list(st), 0, ln-1, k-1)

Switched to an iterative implementation cause python:

def select2(lst, lft, rght, k):
    while True:
        if lft == rght:     #Return the sole element of the list if it is already sorted.
            return lst[lft]
        pvtIdx = lft + int(random()*(rght - lft))       #Generate a random pivot index between `lft` and `rght` (both inclusive).
        pvtIdx = partition(lst, lft, rght, pvtIdx)  #The index of the pivot value in it's sorted position.
        if k == pvtIdx:     #If that index corresponds to the desired index.
            return lst[k]
        elif k < pvtIdx:    #Insert another element to its sorted position in the partition of the list that the desired element resides in.
            right = pvtIdx - 1
            continue
        else:
            left = pvtIdx+1     #Insert another element to its sorted position in the partition of the list that the desired element resides in.
            continue

def nS4(lst = z, k = 16):
    st = set(lst)
    ln = len(st)
    return None if k > ln else select2(list(st), 0, ln-1, k-1)

I tried using a heap:

def nS5(lst = z, k = 16):
    lst = list(set(lst))
    ln = len(lst)
    if k > ln:
        return None
    heapify(lst)
    for i in range(k):
        current = heappop(lst)
    return current

I tried optimising my heap:

def nS6(lst = z, k = 16):
    heapify(lst)
    st = set()
    count = 0
    while count < k:
        if not lst:
            return None
        current = heappop(lst)
        if current not in st:
            st.add(current)
            count += 1
    return current

I tried combining multiple functions together to leverage asymptotics:

def nS7(lst, k):
    if len(lst) < 100000:
        return nS6(lst, k)
    return nS4(lst, k)

I benchmarked my functions (using a fifty element list):

print("t(nS1):\t\t", t(nS1, number = 1000000))
print("t(nS2):\t\t", t(nS2, number = 1000000))
# print("t(nS3):\t\t", t(nS3, number = 1000000))
# print("t(nS4):\t\t", t(nS4, number = 1000000))
print("t(nS5):\t\t", t(nS5, number = 1000000))
print("t(nS6):\t\t", t(nS6, number = 1000000))

(I commented out the two implementations of quickselect because previous benchmarks had shown they took orders of magnitude more time than some of the other implementations).

If others hadn't solved the kata, I would have called bullshit. As it is, I've invested way too much effort into it already.

Answer 1

The list -> set ->list conversion seems pretty wasteful, especially if there are not too many dupes. Besides, it introduces \$O(N)\$ space complexity. Even though set promises a constant time inserts and lookups, the constant could be quite large (for large sets) due to the poor referential locality (hence plenty of cache misses). That said, since you try to sort/heapify an entire list, the time complexity would be at least \$O(N \log N)\$.

The kata doesn't specify how big n could be. If it is comparable to N, there is nothing you could do, besides ditching the set. Just sort the list, and linearly traverse it discarding dupes.

However, if n is much less than N, consider using a fixed-size heap of at most n entries. Bite the bullet and implement sift_down with a couple of twists:

1. If you hit an element equal to the one being sifted, discard the latter.
2. If the element is sifted beyond n, discard it.

Once the entire list is processed, either the heap is not filled completely (return None, or return the maximum.

The time complexity of this approach is \$O(N \log n)\$, and the space complexity is \$O(n)\$.

Answer 2

(I commented out the two implementations of quickselect because previous benchmarks had shown they took orders of magnitude more time than some of the other implementations).

In my testing, for 50-element lists, ns3 is an order of magnitude slower than ns1, and ns4 is an order of magnitude slower than ns3. But I accidentally managed to speed up ns4 by an order of magnitude when renaming the variables to be (IMO) more legible, and on investigating I discovered why:

def select2(lst, lft, rght, k):
    while True:
        if lft == rght:     #Return the sole element of the list if it is already sorted.
            return lst[lft]
        ...
        elif k < pvtIdx:    #Insert another element to its sorted position in the partition of the list that the desired element resides in.
            right = pvtIdx - 1
            ...

The problem with using mangled names is that if you forget to mangle one, it's not so easy to notice.

---

Also, scale up to 5000-element lists and the same two functions become faster than ns1. If you're finding that the code times out for million-element lists, you need to profile with long lists, not really short ones.

Further contributing to the benchmarking flaws, because the benchmarking approach shares the same list between all of the methods, it's not at all a fair test. Where you define z, add

lenz = len(z)

and then add

assert len(arr) == lenz

to the start of all of the nS functions. You'll find that the reason nS6 is so much faster is that it's working on a shorter list.

---

def nS6(lst = z, k = 16):
    ...
    st = set()
    ...
        current = heappop(lst)
        if current not in st:
            st.add(current)
            ...

If the heap is buggy, this function is almost certainly buggy. If it isn't buggy, you don't need st. It suffices to track the previous element popped from the heap and compare to that. I find that this change gives about a 10% to 15% speedup.

PS The same applies to nS2.

---

My benchmark code:

from random import randint, random, seed
from heapq import *
import timeit

t = timeit.timeit
seed(12345)
z0 = [randint(-2147483648, 2147483647) for i in range(0, 500000)]
k0 = len(z0) // 2

def nS1(arr = None, n = k0):
    if arr == None:
        arr = list(z0)

    st = set(arr)
    return sorted(st)[n-1] if n <= len(st) else None

###   and nS2, etc. with similar modifications   ###

# A variant with the change I mention above
def nS2b(arr = None, n = k0):
    if arr == None:
        arr = list(z0)

    prev = None
    count = 0
    for i in sorted(arr):
        if i != prev:
            count += 1
            if count == n:
                return i
            prev = i
    return None

# A variant on nS4 which special-cases when the range gets small.
def nS4b(arr = None, k = k0):
    if arr == None:
        arr = list(z0)

    st = set(arr)
    ln = len(st)
    if k > ln: return None
    arr = list(st)
    left = 0
    right = ln - 1
    k -= 1
    while True:
        if right - left < 10:
            final = sorted(arr[left:right+1])
            return final[k - left]
        pivotIndex = left + int(random()*(right - left))
        pivotIndex = partition(arr, left, right, pivotIndex)
        if k == pivotIndex:
            return arr[k]
        elif k < pivotIndex:
            right = pivotIndex - 1
        else:
            left = pivotIndex + 1


# A variant of nS6 with the change I suggest above
def nS6(arr = None, k = k0):
    if arr == None:
        arr = list(z0)

    heapify(arr)
    count = 0
    prev = None
    while count < k:
        if not arr:
            return None
        current = heappop(arr)
        if current != prev:
            prev = current
            count += 1
    return current


# My own idea for how to speed things up: radix select
def nS8(arr = None, k = k0):
    if arr == None:
        arr = list(z0)

    # Exploit the knowledge that we're working with 32-bit integers
    offset = 2147483648
    arr = [i + offset for i in set(arr)]
    if k > len(arr):
        return None

    shift = 30
    while len(arr) > 1:
        buckets = [[] for i in range(8)]
        for elt in arr:
            buckets[(elt >> shift) & 7].append(elt)
        for bucket in buckets:
            if k <= len(bucket):
                arr = bucket
                break
            else:
                k -= len(bucket)
        shift -= 3
    return arr[0] - offset


# Suggested in comments by Peilonrayz
def Peilonrayz(arr = None, k = k0):
    if arr == None:
        arr = list(z0)

    st = set(arr)
    if k > len(st):
        return None

    return nsmallest(k, st)[-1]


# For benchmarking just list(set(arr))
def uniq(k = k0):
    arr = list(set(z0))
    return arr[k - 1] if k <= len(arr) else None


def test(fn):
    testcases  = [
        ([1, 3, 4, 5], 7, None),
        ([4, 3, 4, 5], 4, None),
        ([45, -10, 4, 5, 4], 4, 45)
    ]
    for testcase in testcases:
        result = fn(testcase[0], testcase[1])
        if result != testcase[2]:
            print(fn.__name__, "failed test case", testcase, "giving", result)
            return float('+inf')

    return t(fn, number = 100)


if __name__ == "__main__":
    implementations = [nS1, nS2, nS2b, nS3, nS4, nS4b,
                       nS5, nS6, nS6b, nS8, Peilonrayz]
    timed = [(test(fn), fn.__name__) for fn in implementations]
    for result in sorted(timed):
        print(result)

    print("---")
    print("t(uniq):\t", t(uniq, number = 100))

Output:

(24.560783660940388, 'nS8')
(27.097620791058496, 'nS2b')
(27.39887558242293, 'nS6b')
(30.668106617453745, 'nS2')
(32.12385269414622, 'nS1')
(32.97220054667446, 'nS6')
(36.23331559560749, 'nS3')
(36.571778446890335, 'nS5')
(37.13606558411453, 'nS4b')
(37.48886835011808, 'nS4')
(108.40215040129226, 'Peilonrayz')
---
t(uniq):         7.5451649473291695