Feature Request: SortedDict key function based on values

[smmalis37]

I would like to be able to have a SortedDict where the sorting is determined by the value being added, or some combination of the key and the value. Basically, the ability for me to provide a key function that takes two arguments.

[grantjenks]

How about this recipe:

import sortedcontainers as sc

class KeyWrapper(object):
    "Wraps a key function and applies with (key, value) item."
    def __init__(self, key, _dict):
        self._key = key
        self._dict = _dict
    def __call__(self, key):
        "Apply key function to key with value."
        return self._key(key, self._dict[key])

class CustomSortedDict(sc.SortedDict):
    "Requires a key function that accepts two arguments: key, value."
    def __init__(self, *args, **kwargs):
        assert len(args) > 0
        args = list(args)
        key = args[0]
        args[0] = KeyWrapper(key, self)
        super(CustomSortedDict, self).__init__(*args, **kwargs)

Example usage:

In [2]: d = CustomSortedDict(lambda k, v: v, a=5, b=4, c=3, d=2, e=1)

In [3]: d
Out[3]: CustomSortedDict(<__main__.KeyWrapper object at 0x10513e910>, 1000, {'e': 1, 'd': 2, 'c': 3, 'b': 4, 'a': 5})

[smmalis37]

When I try to add an element to the CustomSortedDict I get a KeyError on the line

return self._key(key, self._dict[key])

The stack trace above that is line 180 in sorteddict.py and line 65 in sortedlistwithkey.py

[grantjenks]

Hmm, yeah. That's a problem. In that particular case, the dict doesn't yet contain the item when the key function is computed. I'm afraid other parts of the code assume that the key function operates only on the key and not the item or value.

You might try the PriorityDict at https://github.com/grantjenks/prioritydict/blob/master/prioritydict.py which orders elements by value. I haven't published it to PyPI yet but it's pretty far along.

[grantjenks]

Actually, if you override just a couple more methods, I think it works. You need to override the getitem and setitem methods:

class CustomSortedDict(sc.SortedDict):
    "Requires a key function that accepts two arguments: key, value."
    def __init__(self, *args, **kwargs):
        assert len(args) > 0
        args = list(args)
        key = args[0]
        args[0] = KeyWrapper(key, self)
        super(CustomSortedDict, self).__init__(*args, **kwargs)

    def __delitem__(self, key):
        self._list_remove(key)
        self._delitem(key)

    def __setitem__(self, key, value):
        if key in self:
            del self[key]
        self._setitem(key, value)
        self._list_add(key)

[smmalis37]

That definitely fixed the error, however deleting things doesn't seem to be working. That might be due to my code however, I'll keep you posted. Thank you so much for all the help.

[smmalis37]

Yep, that was my bug, it seems like everything is working perfectly now. Thank you so so much. Because I'm using this code in a project, how would you like to be credited and are there any license restrictions on the code you wrote in this issue?

[grantjenks]

Thanks for asking. Same license as for the project. Source available under Apache 2 License: http://www.grantjenks.com/docs/sortedcontainers/#sortedcontainers-license

[grantjenks]

I think to make this recipe work you'll also need to override the copy methods. Here's the whole pattern, updated:

import sortedcontainers as sc

class KeyWrapper(object):
    "Wraps a key function and applies with `(key, value)` item."
    def __init__(self, func, _dict):
        self._func = func
        self._dict = _dict
    def __call__(self, key):
        "Apply key function to `(key, value)` item."
        return self._func(key, self._dict[key])

class CustomSortedDict(sc.SortedDict):
    "Requires a key function that accepts two arguments: key, value."
    def __init__(self, *args, **kwargs):
        assert len(args) > 0 and callable(args[0])
        args = list(args)
        key = self._original_key = args[0]
        args[0] = KeyWrapper(key, self)
        super(CustomSortedDict, self).__init__(*args, **kwargs)

    def __delitem__(self, key):
        """
        Remove ``d[key]`` from *d*.
        Raise KeyError if *key* is not in the dictionary.
        """
        if key not in self:
            raise KeyError(key)
        self._list_remove(key)
        self._delitem(key)

    def __setitem__(self, key, value):
        """Set `d[key]` to *value*."""
        if key in self:
            self._list_remove(key)
            self._delitem(key)
        self._setitem(key, value)
        self._list_add(key)

    def copy(self):
        """Return a shallow copy of the sorted dictionary."""
        return self.__class__(self._original_key, self._load, self.iteritems())

    __copy__ = copy

[smmalis37]

In all the usage I've been making of the previous recipe it's been working perfectly fine, so I have no idea how necessary that is.

[grantjenks]

Well assuming you've never called copy, probably doesn't matter. But I wanted to update this thread because I'm thinking of creating a sortedcontainers.recipes sub-module where I make things like these available.

[grantjenks]

I've published this recipe as part of a new package named sortedcollections. It's still young but I hope to group a number of these recipes together into that package:

• Sorted Collections Documentation
• sortedcollections on PyPI
• sortedcollections on Github