LeetCode Solutions Blog

Design an Expression Tree With Evaluate Function

Number: 1768

Difficulty: Medium

Paid? Yes

Companies: Amazon

Problem Description

Build a binary expression tree from the given postfix tokens where each number is a leaf and each operator is an internal node. The tree must support an evaluate function that computes the result of the represented expression.

Key Insights

Use a stack to build the tree from postfix notation.
When a number is encountered, create a leaf node; when an operator is encountered, pop two nodes as its children.
The evaluate function works recursively: if the node is an operator, evaluate its left and right subtrees and then apply the operation.
For a modular design, store operator functionality in a mapping so that additional operators can be added without modifying evaluate logic.

Space and Time Complexity

Time Complexity: O(n), where n is the number of tokens, since each token is processed once.
Space Complexity: O(n) for storing nodes in the stack and recursion stack for tree evaluation.

Solution

The algorithm processes the postfix tokens one by one using a stack. For each token:

If it is an operand (number), create a node representing a number and push it onto the stack.
If it is an operator, pop the top two nodes from the stack (right operand first, then left) and create a new operator node having these nodes as children. Then push the new node onto the stack. After processing all tokens, the stack contains the root of the expression tree.

The evaluate function is implemented recursively. For a leaf node, it returns the number value directly. For an operator node, it retrieves and applies the operator function (from a mapping), which ensures the design is modular and can be extended to new operators without changes to this function.

Code Solutions

# Node interface and classes for the expression tree.

# Define a base Node class. We provide an evaluate method.
class Node:
    def evaluate(self):
        raise NotImplementedError("Subclasses must implement evaluate()")

# Leaf node representing numbers
class NumNode(Node):
    def __init__(self, value):
        self.value = value
    def evaluate(self):
        # Directly return the stored numeric value
        return self.value

# Operator node representing an arithmetic operation.
class OpNode(Node):
    # mapping of operator symbols to their functions, allowing for modular design.
    operator_map = {
        '+' : lambda x, y: x + y,
        '-' : lambda x, y: x - y,
        '*' : lambda x, y: x * y,
        '/' : lambda x, y: int(x / y) # using int() to simulate truncation towards zero
    }
    
    def __init__(self, op, left, right):
        self.op = op
        self.left = left
        self.right = right
        
    def evaluate(self):
        # Recursively evaluate left and right subtrees
        left_val = self.left.evaluate()
        right_val = self.right.evaluate()
        # Apply the operator function from the mapping
        operation = self.operator_map[self.op]
        return operation(left_val, right_val)

# The main ExpressionTree class that builds the tree from postfix tokens.
class ExpressionTree:
    # The constructor takes postfix tokens and builds the tree.
    def __init__(self, postfix):
        self.root = self.buildTree(postfix)
        
    def buildTree(self, postfix):
        # use a stack to construct the tree based on postfix notation.
        stack = []
        for token in postfix:
            if token in "+-*/":
                # Pop right and left nodes (right first, then left)
                right = stack.pop()
                left = stack.pop()
                # Create a new operator node and push it to the stack.
                node = OpNode(token, left, right)
                stack.append(node)
            else:
                # Create a leaf node for the operand and push it.
                stack.append(NumNode(int(token)))
        # The last element in the stack is the root.
        return stack[-1]
    
    def evaluate(self):
        # Evaluate the expression starting from the root.
        return self.root.evaluate()

# Example usage:
# tokens = ["3", "4", "+", "2", "*", "7", "/"]
# tree = ExpressionTree(tokens)
# print(tree.evaluate())  # Output should be 2