exercises/dsa/ch5_binary_tree/dsa_0509_expression_tree.cpp

#include <iostream>
#include <cassert>
#include <functional>
#include <stack>
#include <vector>
#include <unordered_map>

using Expression = std::string; // 没有数字

bool is_operator(const char c) {
    return c == '+' || c == '-' || c == '*' || c == '/'
        || c == '(' || c == ')';
}

int get_priority(char c) {
    assert(is_operator(c));
    if (c == '+' || c == '-') { return 1; }
    if (c == '*' || c == '/') { return 2; }
    throw std::invalid_argument("Invalid operator!");
}

class ExpressionTree {
public:
    struct Node {
        char value;
        int depth = -1;
        Node *father = nullptr, *left = nullptr, *right = nullptr;

        explicit Node(char c) : value(c) {}

        [[nodiscard]] bool is_operator() const {
            return ::is_operator(value);
        }

        [[nodiscard]] int get_priority() const {
            return ::get_priority(value);
        }
    };

    Node* root = nullptr;
    int depth = -1;

    void init_depths() {
        dfs([this](Node* node, int dfs_depth) {
            node->depth = dfs_depth;
            depth = std::max(depth, dfs_depth);
        });
        depth++;
    }

    void init_with_infix(const Expression& expression) {
        std::stack<Node*> ops;
        std::stack<Node*> nums;

        for (char c : expression) {
            if (!is_operator(c)) {
                nums.push(new Node(c));
                continue;
            }
            if (c == '(') {
                ops.push(new Node(c));
                continue;
            }
            if (c == ')') {
                while (ops.top()->value != '(') {
                    construct_one_op(ops, nums);
                }
                ops.pop();
                continue;
            }
            Node* node = new Node(c);
            while (!ops.empty() && ops.top()->value != '('
                && node->get_priority() <= ops.top()->get_priority()) {
                construct_one_op(ops, nums);
            }
            ops.push(node);
        }
        while (!ops.empty()) {
            construct_one_op(ops, nums);
        }
        root = nums.top();
        init_depths();
    }

protected:
    static void construct_one_op(std::stack<Node*>& ops, std::stack<Node*>& nums) {
        auto op = ops.top();
        ops.pop();
        auto num_right = nums.top();
        nums.pop();
        auto num_left = nums.top();
        nums.pop();
        op->left = num_left;
        op->right = num_right;
        num_left->father = op;
        num_right->father = op;
        nums.push(op);
    }

public:
    [[nodiscard]] Expression get_postfix() const {
        Expression expression;
        // 此代码正确，但是本人想试试 dfs
        // std::stack<std::pair<Node*, short>> stack; // 0: 没干过活，1: 玩过左，2: 玩过右
        // stack.emplace(root, 0);
        // while (!stack.empty()) {
        //     // std::clog << "Top: " << stack.top().first->value << ", flag=" << stack.top().second << std::endl;
        //     if (stack.top().second == 0) {
        //         stack.top().second++;
        //         if (stack.top().first->left != nullptr) {
        //             stack.emplace(stack.top().first->left, 0);
        //         }
        //     } else if (stack.top().second == 1) {
        //         stack.top().second++;
        //         if (stack.top().first->right != nullptr) {
        //             stack.emplace(stack.top().first->right, 0);
        //         }
        //     } else {
        //         expression += stack.top().first->value;
        //         stack.pop();
        //     }
        // }
        dfs([&expression](Node* node, int _) {
            expression.push_back(node->value);
        }, 2);
        return expression;

    }

    std::vector<std::string> get_string_lines() {
        std::vector<std::string> lines(depth * 2);
        dfs_with_id([&lines](Node* node, unsigned long long node_id) {
            lines[node->depth * 2].resize(node_id, ' ');
            lines[node->depth * 2][node_id - 1] = node->value;
            if (node->left != nullptr) {
                lines[node->depth * 2 + 1].resize(node_id - 1, ' ');
                lines[node->depth * 2 + 1][node_id - 2] = '/';
            }
            if (node->right != nullptr) {
                lines[node->depth * 2 + 1].resize(node_id + 1, ' ');
                lines[node->depth * 2 + 1][node_id] = '\\';
            }
        }, 2);
        lines.pop_back();
        return lines;
    }

    long long calculate(const std::unordered_map<char, int>& variable_map, Node* node = nullptr) {
        if (node == nullptr) { node = root; }
        if (node->left == nullptr) { return variable_map.at(node->value); }
        long long left = calculate(variable_map, node->left);
        long long right = calculate(variable_map, node->right);
        if (node->value == '+') { return left + right; }
        if (node->value == '-') { return left - right; }
        if (node->value == '*') { return left * right; }
        if (node->value == '/') { return left / right; }
        throw std::logic_error("");
    }

protected:
    void dfs(const std::function<void(Node*, int)>& func,
             int func_entry = 0,
             int dfs_depth = 0,
             Node* node = nullptr) const {
        if (dfs_depth == 0) { node = root; }
        if (func_entry == 0) { func(node, dfs_depth); }
        if (node->left != nullptr) { dfs(func, func_entry, dfs_depth + 1, node->left); }
        if (func_entry == 1) { func(node, dfs_depth); }
        if (node->right != nullptr) { dfs(func, func_entry, dfs_depth + 1, node->right); }
        if (func_entry == 2) { func(node, dfs_depth); }
    }

    void dfs_with_id(const std::function<void(Node*, unsigned long long)>& func,
                     int func_entry = 0,
                     unsigned long long node_id = 0,
                     Node* node = nullptr) const {
        if (node == nullptr) { node = root, node_id = 1 << (depth - 1); }
        // std::clog << node_id << std::endl;
        unsigned long long depth_bit = 1ull << (depth - node->depth - 1);
        unsigned long long flipped_depth_bit = ~0ull ^ depth_bit;
        if (func_entry == 0) { func(node, node_id); }
        if (node->left != nullptr) {
            dfs_with_id(func, func_entry, (node_id & flipped_depth_bit) | (depth_bit >> 1), node->left);
        }
        if (func_entry == 1) { func(node, node_id); }
        if (node->right != nullptr) {
            dfs_with_id(func, func_entry, node_id | (depth_bit >> 1), node->right);
        }
        if (func_entry == 2) { func(node, node_id); }
    }

    friend std::ostream& operator<<(std::ostream& os, ExpressionTree expression_tree) {
        for (const auto& line : expression_tree.get_string_lines()) {
            os << line << std::endl;
        }
        return os;
    }
};

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    Expression infix_expression;
    std::cin >> infix_expression;

    ExpressionTree expression_tree;
    expression_tree.init_with_infix(infix_expression);
    std::cout << expression_tree.get_postfix() << std::endl;
    std::cout << expression_tree;

    int n;
    std::cin >> n;
    std::unordered_map<char, int> variable_map;
    for (int i = 0; i < n; i++) {
        char c;
        int x;
        std::cin >> c >> x;
        variable_map[c] = x;
    }
    std::cout << expression_tree.calculate(variable_map) << std::endl;

    return 0;
}