Python 3.5 - 'Infect opened tree branches'

Question

I know the title sounds a little weird, but of course this is just an analogy to what I actually need.

So, lets assume I have a tree like this:

A
┃
┣━━ B
┃   ┣━━ D
┃   ┣━━ E
┃   ┃   ┗━━ H
┃   ┗━━ F
┃       ┗━━ I
┗━━ C
    ┗━━ G

with one of the leaves (or branches) infected with some disease.

Traversing the tree will infect all the 'opened' branches/leaves at the traversal time, but not newly opened ones. Lets assume that branch E is infected - traversing the tree yield infected F and C branches, since they already 'opened' in this iteration, but not I and G.

The python code I have so far is (infection_test.py):

#!/usr/bin/env python
from itertools import chain

class Node():
    def __init__(self, name, infected=False):
        self.name = name
        self.children = []

        self.infected = infected

    def __str__(self):
        return 'Node ' + self.name + (' *** INFECTED ***' if self.infected else '')

A = Node('A');B = Node('B');C = Node('C')
D = Node('D');E = Node('E', True);F = Node('F');
G = Node('G');H = Node('H');I = Node('I');

A.children = [B, C]
B.children = [D, E, F]
E.children = [H]
F.children = [I]
C.children = [G]

def traverse_tree(node, level=0):
    print ('    '*level, node)
    level += 1
    infected_found = False
    for child in node.children:
        if child.infected:
            infected_found = True
        traverse_tree(child, level)
        child.infected = infected_found

print('First traversal:')
traverse_tree(A)
print('\nAfter Infection:')
traverse_tree(A)

Which outputs:

First traversal:
 Node A
     Node B
         Node D
         Node E *** INFECTED ***
             Node H
         Node F
             Node I
     Node C
         Node G

After Infection:
 Node A
     Node B
         Node D
         Node E *** INFECTED ***
             Node H
         Node F *** INFECTED ***
             Node I
     Node C
         Node G

How can I make 'higher level' branches (like C) to be infected, without influence on the next iterations of traverse_tree?

(I hope that 'opened branches' is clear enough, but just to make sure it is - those are the branches that yields from the for child loop already, when the infected branch discovered)

Answer 1

OK, so I found a solution, but I think I'll still wait for a better one.

I just edited the original Node and added the p_infected attribute. This will help me to mark all the branches that I need to infect later on. As soon as I found some infected branch - it infects all of his parents up to the root. Then, I traverse the tree, infect back the children of those branches, and remove the 'parent infection'.

The current code look like this:

from itertools import chain

class Node():
    def __init__(self, name, infected=False, p_infected=False):
        ...
        self.parent   = None
        self.p_infected = p_infected

    def add_children(self, childs = []):
        self.children.extend(childs)
        for node in childs:
            node.parent = self

    def __str__(self):
        return 'Node ' + self.name + \
        (' ***INFECTED***' if self.infected else '') + \
        (' ***P_INFECTED***' if self.p_infected else '')

A = Node('A');B = Node('B');C = Node('C')
D = Node('D',True);E = Node('E');F = Node('F');
G = Node('G');H = Node('H');I = Node('I');

A.add_children([B,H])
B.add_children([C, D, F])
D.add_children([E])
F.add_children([G])
H.add_children([I])

def infect_parent(node):
    if node.parent:
        node.parent.p_infected = True
        infect_parent(node.parent)

def traverse_tree(node, level=0):
    print ('    '*level + str(node))
    if node.infected:
        node.p_infected = True
        infect_parent(node)
    level += 1
    for child in node.children:
        traverse_tree(child, level)


def infect_children(node):
    infect_flag = False
    for child in node.children:
        if infect_flag:
            child.infected = True
        infect_flag = child.p_infected
        infect_children(child)

def remove_parent_infection(node):
    node.p_infected = False
    for child in node.children:
        remove_parent_infection(child)

print('First travesal:')
traverse_tree(A)
infect_children(A)
remove_parent_infection(A)
print('\nAfter infection:')
traverse_tree(A)

And the output, as desired:

First travesal:
Node A
    Node B
        Node C
        Node D ***INFECTED***
            Node E
        Node F
            Node G
    Node H
        Node I

After infection:
Node A
    Node B
        Node C
        Node D ***INFECTED***
            Node E
        Node F ***INFECTED***
            Node G
    Node H ***INFECTED***
        Node I

But, as I wrote before, I'm still open to a better solution. I have a hunch that the above could be done with just the original 'traversal' function, and if someone could find how - I'll be happy..

Thanks.