python – How to find related nodes in a tree (not binary)?

Question:

at the moment I want to understand how to find all connected nodes in a tree.

Here is an example code:

class Tree:

    def __init__(self, root, subtrees=None):
        self.root = root
        self.subtrees = subtrees[:] if subtrees else []


    def con(self, v1, v2):
        # base cases
        if v1 and v2 is None:
            return []
        elif self.root is None:
            return []
        elif v1 == v2:
            return [v1]
        else:
            path = []
            for subtree in self.subtrees:
                path += subtree.con(v1, v2)
            return path

How to recurse through all elements of a tree? Because in my code it only goes through the branches, and because of this, the answer is in the form of an empty list. I can't figure out how to find the nodes along with the root.

Output should be something like:

t1 = Tree(2, [Tree(5)])

t2 = Tree(4, [Tree(6), Tree(7)])

t = Tree(1, [t1, t2])

print(t.con(5,4))

5, 2, 1, 4

Answer:

class Tree:

    def __init__(self, root, subtrees=None):
        self.root = root
        self.subtrees = subtrees[:] if subtrees else []

    def con(self, v1, v2, path = []):
        if (v1 is None and v2 is None) or (self.root is None):
            return []
        elif v1 == v2:
            return [v1]
        elif self.root != v1 and self.root != v2:
            path.append(self.root);
            for sub in self.subtrees:
                sub.con(v1, v2, path)
        elif self.root == v1:
            path.append(v1)
        elif self.root == v2:
            path.append(v2)
        return path

Well, here's something like this (sorry, I didn't understand whether the order of the elements is important to you, if anything, you can think about how to redo it).

The bottom line is this: we start recursion from the root of the tree and go through all subtrees, start recursion from them until we meet the vertex we need, simultaneously writing all the vertices to path .

This approach has a clear disadvantage: t.con(5, 1) will simply return 1

So let's change something:

class Tree:

    parent = None

    def __init__(self, root, subtrees=None):
        self.root = root
        self.subtrees = subtrees[:] if subtrees else []
        if subtrees:
            for sub in subtrees:
                sub.parent = self;

    def find(self, path = []):
        if self.root is None:
            return path
        else:
            if not (self.root in path): path.append(self.root)
            if self.parent: self.parent.find(path)

    def con(self, v1, v2, path = []):
        if (v1 is None and v2 is None) or (self.root is None):
            return []
        elif v1 == v2:
            return [v1]
        elif self.root != v1 or self.root != v2:
            for sub in self.subtrees:
                sub.con(v1, v2, path)
        if self.root == v1:
            self.find(path)
        elif self.root == v2:
            self.find(path)
        return path

We will store the ancestor of each subtree. In the recursion, we first go down to the desired roots, and from them we launch another recursion to the very top, and it is the second recursion that will write path .

But there is a problem here too: t.con(5, 2) will output [5, 2, 1] instead [5, 2] Therefore, I wrote this crutch:

class Tree:

    parent = None

    def __init__(self, root, subtrees=None):
        self.root = root
        self.subtrees = subtrees[:] if subtrees else []
        if subtrees:
            for sub in subtrees:
                sub.parent = self;

    def find(self, v1, v2, path = [], first = False):
        if self.root is None:
            return path
        else:
            if not (self.root in path): 
                path.append(self.root)
            else:
                return path
            if self.parent and (self.root != v1 and self.root != v2 or first): 
                self.parent.find(v1, v2, path)

    def con(self, v1, v2, path = [], par = False):
        if (v1 is None and v2 is None) or (self.root is None):
            return []
        elif v1 == v2:
            return [v1]

        if self.root == v1:
            if par: 
                self.find(v1, v2, path, True)
                if self.root != v1 or self.root != v2:
                    for sub in self.subtrees:
                        sub.con(v1, v2, path, True)
            else:
                if self.root != v1 or self.root != v2:
                    for sub in self.subtrees:
                        sub.con(v1, v2, path, False)
                self.find(v1, v2, path, True)

        elif self.root == v2:
            if par: 
                self.find(v1, v2, path, True)
                if self.root != v1 or self.root != v2:
                    for sub in self.subtrees:
                        sub.con(v1, v2, path, True)
            else:
                if self.root != v1 or self.root != v2:
                    for sub in self.subtrees:
                        sub.con(v1, v2, path, False)
                self.find(v1, v2, path, True)

        elif self.root != v1 or self.root != v2:
            for sub in self.subtrees:
                sub.con(v1, v2, path, False)


        return path

It's just that I need to run find from a vertex deeper than the second one (so we don't have to go too high to the root of the source tree) and I have no idea how to do it without such a crutch or a global variable.

Maybe you can leave the crutch, but somehow mix the elements of the code and it will turn out more compact, but I've been writing this for almost 2 hours, maybe my eyes are blurry, sorry. I would be glad if someone corrects me.

Scroll to Top