Set operations in Python

Alex cut_me_out at hotmail.com
Wed Sep 20 19:05:54 EDT 2000


> Is there a good way to do set operations?  I'm currently writing my
> own routines to do a union, intersection, etc. on two or more lists,
> but I figure this has got to be common enough that someone has
> probably done it first.  I didn't find anything on the python.org Web
> site, though, so I thought I'd ask here just in case.

Here's what I use.

Alex.

from types import ListType, TupleType

class Set:

    def __init__(self, *args):
        self.set = {}
        # Check whether the elements
        if(len(args) == 1) and(type(args[0]) in(ListType, TupleType)):
            args = args[0]
        # Add the elements to the set.    
        for elt in args:
            self.set[elt] = None

    def __and__(self, other):
        '''Take the intersection of two sets.'''
        # Find the shorter of the pair to iterate over.
        if len(other.set) < len(self.set):
            set_one = other.set
            set_two = self.set
        else:
            set_one = self.set
            set_two = other.set
        set_two_has_key = set_two.has_key
        intersection = {}
        for elt in set_one.keys():
            if set_two_has_key(elt):
                intersection[elt] = None
        new_set = Set()
        new_set.set = intersection
        return new_set

    def __len__(self):
        return len(self.set)

    def incorporate(self, other):
        '''Add the elements of other to self.'''
        self.set.update(other.set)

    def add(self, elt):
        self.set[elt] = None

    def remove(self, elt):
        del self.set[elt]

    def __delitem__(self, elt):
        self.remove(elt)

    def incorporate_sequence(self, seq):
        for elt in seq:
            self.set[elt] = None

    def copy(self):
        '''Return a copy of self.'''
        new_set = Set()
        new_set.set = self.set.copy()
        return new_set

    def subset(self, other):
        '''Return whether argument is a subset of self.'''
        other_set_has_key = other.set.has_key
        for elt in self.set.keys():
            if not other_set_has_key(elt):
                break
        else: # Didn't break out, all elts found
            return 1
        return 0

    def __add__(self, other):
        '''Return union of self with other.'''
        if len(other.set) < len(self.set):
            set_one = other
            set_two = self
        else:
            set_one = self
            set_two = other
        new_set = set_two.copy()
        new_set.incorporate(set_one)
        return new_set

    def values(self):
        '''Return elements as a list.'''
        return self.set.keys()

    def image(self, function):
        '''Returns the set of return values for function applied
        to the set'''
        return Set(map(function, self.set.keys()))

    def __sub__(self, other):
        '''Return the difference of self and other.'''
        other_set_has_key = other.set.has_key
        new_set = Set()
        new_set_set = new_set.set
        for elt in self.set.keys():
            if not other_set_has_key(elt):
                new_set_set[elt] = None
        return new_set

    def __contains__(self, elt):
        '''Only works in bleeding edge python - 2.0b'''
        return self.set.has_key(elt)

    def member(self, elt):
        return self.__contains__(elt)

    def __repr__(self):
        return 'Set (%s)' % self.values()

    def __cmp__(self, other):
        if hasattr(other, 'set'):
            return cmp(self.set, other.set)
        else:
            return 1

if __name__ == '__main__':
    t = Set()
    print t
    t.incorporate_sequence(range(5))
    print t
    u = Set(range(2, 4))
    s = t - u
    print s
    assert s.subset(t)
    print '%s is a subset of %s' % (s, t)
    assert s & t == s
    print '%s & %s == %s' % (s, t, s)
    assert s + u == t
    print '%s + %s == %s' % (s, u, t)


-- 
Speak softly but carry a big carrot.



More information about the Python-list mailing list