C++ Program for Tree Iterator

Problem

Given n-ary (n children per node) tree, write a C++ iterator to iterate over all the nodes.

Solution

Iterator Interface

One of the potential interface is as follows:

/**
 * If there are more children to be traversed in the current layer
 * @param index
 * @return True if there are children yet to be traversed, false otherwise
 */
virtual bool hasMoreChildren(int index);

/**
 * Move to child at index in a layer
 * @param index
 */
virtual void moveToChildAt(int index);

/**
 * Print the current node
 * @param index
 */
virtual void printNode();

Implementation

For implementing the iterator, we will use a stack to keep track of the node being traversed in the tree.

private:
    Node* mCurrentNode = nullptr;
    std::stack<Node*> mNodeStack;
};

The mCurrentNode and mNodeStack are used in the implementation as follows:

C++ Code

bool TreeIterator::hasMoreChildren(int index) {
    if (mNodeStack.empty()) return false;
    mCurrentNode = mNodeStack.top();
    // If no children, return false.
    if (mCurrentNode->children().empty()) {
        mNodeStack.pop();
        if (!mNodeStack.empty()) {
            mCurrentNode = mNodeStack.top();
        }
        return false;
    }
    auto it = mCurrentNode->children().begin();
    it += index;
    // If all children in this layer are traversed, return false.
    if (it == mCurrentNode->children().end()) {
        mNodeStack.pop();
        if (!mNodeStack.empty()) {
            mCurrentNode = mNodeStack.top();
        }
        return false;
    }
    // There are more children to be drawn, return true
    return true;
}

void TreeIterator::moveToChildAt(int index) {
    if (mCurrentNode->children().empty()) {
        // No children, nothing to do in current layer
        return;
    }
    auto it = mCurrentNode->children().begin();
    it += index;
    if (it == mCurrentNode->children().end()) {
        // All children visited once, nothing to do in current layer
        return;
    }
    mNodeStack.push(*it);
    mCurrentNode = mNodeStack.top();
}

void TreeIterator::printNode() {
    cout << mCurrentNode->data() << endl;
}

Using the iterator

/**
 * Recursively traverse the tree hierarchy
 */
void traverse() {
    int i = 0;
    while (hasChildren(i)) {
        moveToChildAt(i);
        traverse();
        i++;
    }
}

int main() {
    TreeIterator *it;
    it->traverse();
}

Explanation

The above code is a recursively traverses the children of each node in the tree. The stack maintains the current node at the top and the index passed from the traverse() method determines how many children of a particular node have been visited.

Alternate Solutions

Alternate solutions may use the following approach.

Visitor Pattern

Yet to write a working code, but this approach would mark a node as visited once it has been traversed to keep track of which child of a particular node should be visited next.

Parent Tracking

This requires modifying the tree node to contain a pointer to the parent node so that the mCurrentNode can be moved to the parent when all the children of a particular node are visited.