"""PartitionRefinement.py

Maintain and refine a partition of a set of items into subsets,
as used e.g. in Hopcroft's DFA minimization algorithm,
modular decomposition of graphs, etc.

D. Eppstein, November 2003.
"""

class PartitionError(Exception): pass

from sets import Set, ImmutableSet

class PartitionRefinement:
    """Maintain and refine a partition of a set of items into subsets.
    Space usage for a partition of n items is O(n), and each refine
    operation takes time proportional to the size of its argument.
    """

    def __init__(self,items):
        """Create a new partition refinement data structure for the given
        items.  Initially, all items belong to the same subset.
        """
        S = Set(items)
        self._sets = [S]
        self._setids = Set([id(S)])
        self._partition = dict([(x,S) for x in S])
        
    def __getitem__(self,element):
        """Return the set that contains the given element."""
        return self._partition[element]
        
    def __iter__(self):
        """Loop through the sets in the partition."""
        return iter(self._sets)
        
    def __len__(self):
        """Return the number of sets in the partition."""
        return len(self._sets)

    def add(self,element,set):
        """Add a new element to the given partition subset."""
        if self._setids is None:
            raise PartitionError("Partition has been frozen")
        if id(set) not in self._setids:
            raise PartitionError("Set does not belong to the partition")
        if element in self._partition:
            raise PartitionError("Element already belongs to the partition")
        self._partition[element] = set
        set.add(element)

    def remove(self,element):
        """Remove the given element from its partition subset."""
        if self._setids is None:
            raise PartitionError("Partition has been frozen")
        set = self._partition[element]
        del self._partition[element]
        set.remove(element)

    def refine(self,S):
        """Refine each set A in the partition to the two sets
        A & S, A - S.  Return a list of pairs (A & S, A - S)
        for each changed set.  Within each pair, A & S will be
        a newly created set, while A - S will be a modified
        version of an existing set in the partition.
        """
        if self._setids is None:
            raise PartitionError("Partition has been frozen")
        hit = []
        ids = {}
        S = Set(S)
        for x in S:
            if x in self._partition:
                Ax = self._partition[x]
                if id(Ax) not in ids:
                    hit.append(Ax)
                    ids[id(Ax)] = Set()
                ids[id(Ax)].add(x)
        output = []
        for A in hit:
            AS = ids[id(A)]
            if AS != A:
                self._sets.append(AS)
                self._setids.add(id(AS))
                for x in AS:
                    self._partition[x] = AS
                A -= AS
                output.append((AS,A))
        return output

    def freeze(self):
        """Make all sets in S immutable."""
        isets = []
        for S in self:
            I = ImmutableSet(S)
            for x in I:
                self._partition[x] = I
            isets.append(I)
        self._sets = isets
        self._setids = None

@@ -211,16 +211,52 @@
 3.%0A%22%22%22%0A%0A
+from sets import Set, ImmutableSet%0A%0A
 class Pa
@@ -290,44 +290,8 @@
 ss%0A%0A
-from sets import Set, ImmutableSet%0A%0A
 clas
@@ -756,47 +756,17 @@
 s = 
-%5BS%5D%0A        self._setids = Set(%5B
+%7B
 id(S)
-%5D)
+:S%7D
 %0A   
@@ -1065,13 +1065,8 @@
 urn 
-iter(
 self
@@ -1071,16 +1071,28 @@
 lf._sets
+.itervalues(
 )%0A      
@@ -1305,103 +1305,8 @@
 %22%22%22%0A
-        if self._setids is None:%0A            raise PartitionError(%22Partition has been frozen%22)%0A
     
@@ -1336,18 +1336,16 @@
 elf._set
-id
 s:%0A     
@@ -1520,32 +1520,57 @@
 the partition%22)%0A
+        set.add(element)%0A
         self._pa
@@ -1595,33 +1595,8 @@
  set
-%0A        set.add(element)
 %0A%0A  
@@ -1700,109 +1700,8 @@
     
- if self._setids is None:%0A            raise PartitionError(%22Partition has been frozen%22)%0A        set =
  sel
@@ -1713,32 +1713,48 @@
 rtition%5Belement%5D
+.remove(element)
 %0A        del sel
@@ -1778,36 +1778,8 @@
 ent%5D
-%0A        set.remove(element)
 %0A%0A  
@@ -2123,147 +2123,36 @@
     
-if self._setids is None:%0A            raise PartitionError(%22Partition has been frozen%22)%0A        hit = %5B%5D%0A        ids = %7B%7D%0A        S = Set(S)
+hit = %7B%7D%0A        output = %5B%5D
 %0A   
@@ -2265,159 +2265,43 @@
     
-if id(Ax) not in ids:%0A                    hit.append(Ax)%0A                    ids%5B
+hit.setdefault(
 id(Ax)
-%5D = 
+,
 Set()
-%0A                ids%5Bid(Ax)%5D.add(x)%0A        output = %5B%5D
+).add(x)
 %0A   
@@ -2314,15 +2314,26 @@
 or A
+,AS
  in hit
+.items()
 :%0A  
@@ -2347,21 +2347,23 @@
    A
-S
  = 
-ids%5Bid(A)
+self._sets%5BA
 %5D%0A  
@@ -2414,60 +2414,21 @@
 sets
-.append(AS)%0A                self._setids.add(id(AS))
+%5Bid(AS)%5D = AS
 %0A   
@@ -2660,40 +2660,36 @@
     
-isets = %5B%5D%0A        for S in self
+for S in self._sets.values()
 :%0A  
@@ -2797,58 +2797,45 @@
     
-isets.append(I)%0A        self._sets = isets%0A
+self._sets%5Bid(I)%5D = I%0A     
        
+del
  sel
@@ -2844,15 +2844,13 @@
 _set
-ids = None
+s%5Bid(S)%5D
 %0A

Clean up PartitionRefinement.