python: retrieving ceiling key and floor key in a dictionary or a set

Question

I have a dictionary with a bunch of integer keys. for keys I don't have, I'd like to be able to retrieve the smallest and largest keys right before and after they key I want to retrieve but that does not exist.
The Treemap class in java has two methods doing exactly this: ceilingkey() and floorkey().

How can I do this with python?

As an example I have a dictionary like this:

 { 1: "1", 4: "4", 6: "6" ..., 100: "100" } 

If I ask for key 1, I'll retrieve "1", but if I look for key 3, I should get KeyError and hence be able to get floor(3) = 1 and ceil(3) = 4.

Answer 1

def floor_key(d, key):
    if key in d:
        return key
    return max(k for k in d if k < key)

def ceil_key(d, key):
    if key in d:
        return key
    return min(k for k in d if k > key)

I'm not sure how you want to handle border conditions. Note that this will raise an exception (ValueError) if you are asking for the floor/ceiling of a key that's lower/higher than anything in the dict.

Answer 2

You can use the bisect module here, in case the key the not found then it can find the ceil and floor values in O(log N) time.:

>>> import bisect
>>> from random import randint
def get_value(dic, key):
    if key in dic:
        return dic[key]
    else:
        ind = bisect.bisect(keys, key)
        d = {}
        if ind > 0:
            d["floor"] = dic[keys[ind-1]]
        if ind < len(keys):
            d["ceil"] = dic[keys[ind]]
        return d
...     
>>> dic = {randint(0,100) : x  for x in xrange(10)}
>>> dic
{65: 6, 4: 5, 1: 7, 40: 8, 10: 4, 50: 0, 68: 2, 27: 9, 61: 3}

Create a sorted list of keys for bisect:

>>> keys = sorted(dic)
>>> keys
[1, 4, 10, 27, 40, 50, 61, 65, 68]

Now use the function:

>>> get_value(dic, 4)
5

For 3 both ceil and floor are avaiable:

>>> get_value(dic, 3)
{'ceil': 5, 'floor': 7}

0 is smaller than the smallest key, so this will return only ceil:

>>> get_value(dic, 0)
{'ceil': 7}

For keys greater than the largest key only floor value will be returned:

>>> get_value(dic, 70)
{'floor': 2}

Answer 3

Here is a fun one using bisect.

import bisect

def ceiling_key(d, key):
    keys = sorted(d.keys())
    idx = bisect.bisect_left(keys, key)
    if key in d:
        idx += 1
    return keys[idx]

def floor_key(d, key):
    keys = sorted(d.keys())
    idx = bisect.bisect_left(keys, key) - 1
    return keys[idx]


...

>>> ceiling_key(d, 1)
4
>>> ceiling_key(d, 2)
4
>>> ceiling_key(d, 4)
6
>>> ceiling_key(d, 3)
4
>>> ceiling_key(d, 2)
4
>>> ceiling_key(d, 1)
4
>>> ceiling_key(d, 0)
1
>>> ceiling_key(d, 6)
100
>>> floor_key(d, 6)
4
>>> floor_key(d, 7)
6
>>> floor_key(d, 101)
100
>>> floor_key(d, 100)
6