.

wcfs/internal/xbtree.py

@@ -241,7 +241,11 @@ def Restructure(ztree, newStructure):
             assert isinstance(rzn.node, (ztreeType, zbucketType))
         for rn in RNv:
             assert isinstance(rn, _NodeInRange)
-            assert isinstance(rn.node, (Tree|Bucket)
+            assert isinstance(rn.node, (Tree|Bucket))
+
+            assert not hasattr(rn.node, 'Z')
+            rn.node.Z = None
+
 
         # prepare cost matrix
         C = np.zeros((len(RNv), len(RZv)))
@@ -254,6 +258,22 @@ def Restructure(ztree, newStructure):
         # property of D function so that if B2 > B1 (all keys in B2 > all keys
         # in B1) and A < B1.hi, then D(B2, A) > D(B1, A).
         jv, iv = scipy.optimize.linear_sum_assignment(C)
+        for (j,i) in zip(jv, iv):
+            RNv[j].node.Z = RZv.node
+
+        # if N(old) > N(new) - some old nodes won't be linked to  (it is ok)
+        # if N(old) < N(new) - some new nodes won't be linked from - link them
+        # to newly created Z nodes.
+        for j2 in range(len(RNv)):
+            if j2 not in jv:
+                n = RNv[j2].node
+                assert n.Z is None
+                assert isinstance(n, (Tree,Bucket))
+                if isinstance(n, Tree):
+                    zn = ztreeType()
+                else:
+                    zn = zbucketType()
+                n.Z = zn
 
         if len(Av) == len(Bv):
             # XXX assert the result is 1-1 mapping