LinkCutTree Implementation where the tree paths are Python tuples.

linkcuttree.py

__author__ = 'weigla'
__date__ = '2012-12-12'

class LinkCutTree(object):
    def __init__(self):
        self.paths = PathSet()
        self.successor = dict()

    def maketree(self, id):
        p = self.paths.makepath(id)
        self.successor[p] = None

    def link(self, v, w):
        a,b = self._expose(v), self._expose(w)
        p = self.paths.join(b,None,a)
        self.successor[p ] = None

    def cut(self, id):
        self._expose(id)
        p, q = self.paths.split(id)
        self.successor[p] = None
        self.successor[q] = None

    def _expose(self, v):
        exposed = v
        p = None
        while v is not None:
            w = self.successor[self.paths.findpath(v)]
            [q, r] = self.paths.split(v)
            if q is not None:
                self.successor[q] = v
            p = self.paths.join(p, v, r)
            v = w

        self.successor[p] = None

        return self.paths.findpath(exposed)


class PathSet(object):
    def __init__(self):
        self._paths = set()

    def makepath(self, v):
        p = (v,)
        self._paths.add(p)
        return p

    def findpath(self, v):
        for p in self._paths:
            if v in p:
                return p

    def join(self, p, v, q):
        #print(p,v,q)

        if p is None: p = tuple()
        if q is None: q = tuple()

        if v is None:
            r = p+q
        else:
            r = p + (v,) + q

        for val in (p, (v,), q):
            self._paths.discard(val)

        self._paths.add(r)
        return r

    def split(self, v):
        r = self.findpath(v)
        i = r.index(v)

        self._paths.discard(r)

        p, q = tuple(r[:i]), tuple(r[i+1:])

        self._paths.add(p)
        self._paths.add(q)
        self._paths.add((v,))
        return t(p), t(q)


def t( tpl ):
    return tpl if len(tpl) >= 1 else None

test.py

__author__ = 'weigla'

import sys
sys.path.append('../src')

import unittest
from linkcuttree.path import LinkCutTree, PathSet


class PathSetTest(unittest.TestCase):
    def setUp(self):
        self.ps = PathSet()
        for a in "abcd":
            self.ps.makepath(a)


    def test_makepath(self):
        print(self.ps._paths)

    def test_findpath(self):
        for a in "abcd":
            self.ps.findpath(a)

    def test_join(self):
        self.ps.join(
            self.ps.findpath("a"), "b", None)

        self.ps.join(self.ps.findpath("a"),
                     "c", self.ps.findpath("d"))

        self.assertEqual(('a', 'b', 'c', 'd'),self.ps.findpath("a"))


    def test_split(self):
        self.ps.join(
            self.ps.findpath("a"), "b", None)

        self.ps.join(self.ps.findpath("a"),
                     "c", self.ps.findpath("d"))

        self.ps.split("b")
        self.ps.split("c")
        self.ps.split("d")

        print(self.ps._paths)

class LinkCutTreeTest(unittest.TestCase):
    def setUp(self):
        pass

    def test_simple(self):
        lct = LinkCutTree()

        lct.maketree(1)
        lct.maketree(2)
        lct.maketree(3)
        lct.maketree(4)
        lct.maketree(5)

        lct.link(1,2)
        lct.link(1,3)
        lct.link(3,4)
        lct.link(3,5)

        lct.cut(3)


    def test_complex(self):
        lct = LinkCutTree()

        z = globals()
        for c in 'abcdefghijklmnop':
            lct.maketree(c)
            z[c] = c


        lct.link(b, a)
        lct.link(e, b)
        lct.link(h, e)
        lct.link(f, b)
        lct.link(c, a)
        lct.link(d, a)
        lct.link(g, d)
        lct.link(i, g)
        lct.link(j, g)
        lct.link(l, j)
        lct.link(k, g)

        lct.link(n, m)
        lct.link(p, n)
        lct.link(o, m)

        print(lct.paths._paths)
        print(lct.successor)


if __name__ == "__main__":
    unittest.main()