Solve postfixed quadratic equations

postfix-quadratic.cpp

#include <iostream>
#include <vector>
#include <stack>
#include <string>
#include <cmath>
#include <stdexcept>

/**
 * The idea is to turn the quadratic function into a two variable function,
 * where two coefficients have two separate variables. For example,
 * i.e. turn f(x) = ax^2 + bx + c into f(x,y) = ax^2 + by + c
 * Now finding a, b and c coefficients becomes extremely easy.
 * If we set x = 0 and y = 0 and evaluate f(x,y), we get the coefficient c.
 * If we set x = 1 and y = 0 and evaluate f(x,y), we get the value a + c.
 * If we set x = 0 and y = 1 and evaluate f(x,y), we get the value b + c.
 * In this way we can easily handle any form of the quadratic equation.
 *
 * In the case of product form, i.e. f(x) = (x + p) * (x + q) + r,
 * we can just the scan the postfix expression from left to right
 * and get the p, q, r, values; because in this case the equation follows
 * a strict form. Now from p, q, r values, we can easily get all the 
 * quadratic coefficients:
 *
 * a = 1.0
 * b = p + q
 * c = pq + r
 */

using namespace std;

// Print function for a vector
template<typename T>
ostream& operator<<(ostream& os, const vector<T>& v) {
    os << "[";
    if(!v.empty()) {
        size_t n = v.size();
        for(size_t i = 0 ; i < n-1; i++) os << v[i] << ",";
        os << v[n-1];
    }
    os << "]";
    return os;
}

// Print function for a stack
template<typename T>
ostream& operator<<(ostream& os, const stack<T>& s) {
    os << "[";
    if(!s.empty()) {
        vector<T> v;
        stack<T> t = s;
        while(!t.empty()) { v.push_back(t.top()); t.pop(); }
        size_t n = v.size();
        for(size_t i = 0 ; i < n-1; i++) os << v[i] << ",";
        os << v[n-1];
    }
    os << "]";
    return os;
}

// Check if a token is an operator
bool isOp(string token) {
    return token == "+" or token == "-" or token == "*" or token == "/";
}

// Basic postfix evaluation with variable x and y
float eval(vector<string>& expr, float x, float y) {
    stack<float> s;
    for(auto token : expr) {
        if(!isOp(token)) {
            if(token == "x") s.push(x);
            else if(token == "y") s.push(y);
            else s.push(stof(token));
        }
        else {
            float b = s.top(); s.pop();
            float a = s.top(); s.pop();
            if(token == "+") s.push(a + b);
            else if (token == "-") s.push(a - b);
            else if (token == "*") s.push(a * b);
            else if (token == "/") s.push(a / b);
            else throw runtime_error("Invalid operator");
        }
    }
    return s.top();
}

// Turn f(x) into f(x,y)
vector<string> renameVariables(vector<string>& expr) {
    vector<int> xpos;
    for(size_t i = 0 ; i < expr.size() ; i++)
        if(expr[i] == "x") xpos.push_back(i);
    // Only rename the variable for the second term
    // i.e. ax^2 + bx + c becomes ax^2 + by + c
    if(xpos.size() == 3) expr[xpos[2]] = "y";
    // and bx + c becomes by + c
    else if (xpos.size() == 1) expr[xpos[0]] = "y";
    return expr;
}

// Find a, b, and c coefficients from f(x,y)
vector<float> getQuadFormCoeffs(vector<string>& expr) {
    float c = eval(expr, 0.0, 0.0); // if x = 0 and y = 0, we get c
    float a = eval(expr, 1.0, 0.0); a = a - c; // if x = 1 and y = 0, we get a = a - c
    float b = eval(expr, 0.0, 1.0); b = b - c; // if x = 0 and y = 1, we get b = b - c
    return {a, b, c};
}

// Function to detect if a postfix expression is in product form
bool isProductForm(vector<string>& expr) {
    int countX = 0;
    int countProd = 0;
    for(size_t i = 0 ; i < expr.size() ; i++) {
        if(expr[i] == "x") countX++;
        if(expr[i] == "*") countProd++;
    }
    return (countX == 2 and countProd == 1);
}

// Remove all subtraction op, turn (x - p) * (x - q) - r into (x + (-p))*((x + (-q)) + (-r)
vector<string> canonicalize(vector<string>& expr) {
    size_t n = expr.size();
    if(expr[2] == "-") { expr[2] = "+"; expr[1] = "-" + expr[1]; }
    if(expr[5] == "-") { expr[5] = "+"; expr[4] = "-" + expr[4]; }
    if(expr[n-1] == "-") { expr[n-1] = "+"; expr[n-2] = "-" + expr[n-2]; }
    return expr;
}

// Just scan the product form quadratic equation from left to right
// and gather p, q, r values
vector<float> getProdFormCoeffs(vector<string>& expr) {
    float p = 0.0, q = 0.0, r = 0.0;
    p = stof(expr[1]);
    q = stof(expr[4]);
    if(expr.size() == 9) r = stof(expr[7]);
    return {1.0, (p + q), p * q + r};
}

// Find roots
vector<float> roots(float a, float b, float c) {
    if(a > 0.0 or a < 0.0) {
        float d = b * b - 4.0 * a * c;
        float x1 = (-b + sqrt(d)) / (2.0 * a);
        float x2 = (-b - sqrt(d)) / (2.0 * a);
        return {x1, x2};
    }
    else return {(-c / b)};
}

// main function
int main() {
    vector<vector<string>> quadratics = {
        // Quadtratic forms:
        // ax^2 + bx + c (a = 2, b = 9, c = 3)
        {"2", "x", "*", "x", "*", "9", "x", "*", "+", "3", "+"},
        // x^2 + bx + c (a = 1, b = 9, c = 3)
        {"x", "x", "*", "7", "x", "*", "+", "3", "+"},
        // x^2 + x + c (a = 1, b = 1, c = 1), no real roots
        {"x", "x", "*", "x", "+", "1", "+"},
        // ax^2 + c (a = 2, b = 0, c = -2)
        {"2", "x", "*", "x", "*", "-2", "+"},
        // bx + c (a = 0, b = 9, c = -2)
        {"9", "x", "*", "-2", "+"},
        // x + c (a = 0, b = 1, c = -1)
        {"x", "-1", "+"},

        // Product forms (two x's but one '*'):
        // (x + p) * (x + q) + r
        {"x", "3", "+", "x", "5", "+", "*", "2", "+"},
        // (x - p) * (x + q) + r
        {"x", "3", "-", "x", "5", "+", "*", "2", "+"},
        // (x + p) * (x - q) + r
        {"x", "3", "+", "x", "5", "-", "*", "2", "+"},
        // (x - p) * (x - q) + r
        {"x", "3", "-", "x", "5", "-", "*", "2", "+"},
        // (x - p) * (x - q) - r
        {"x", "3", "-", "x", "5", "-", "*", "2", "-"},
        // (x + p) * (x - q)
        {"x", "3", "+", "x", "5", "-", "*"}
    };

    for(auto fx : quadratics) {
        vector<float> coeffs;
        cout << "f(x) = " << fx << endl;
        if(!isProductForm(fx)) {
            vector<string> fxy = renameVariables(fx);
            cout << "f(x,y) = " << fxy << endl;
            coeffs = getQuadFormCoeffs(fxy);
        }
        else {
            vector<string> gx = removeSubOp(fx);
            cout << "g(x) = " << gx << endl;
            coeffs = getProdFormCoeffs(gx);
        }

        cout << "{a, b, c} = " << coeffs << endl;
        if(!coeffs.empty())
            cout << "roots = " << roots(coeffs[0], coeffs[1], coeffs[2]) << endl;
    }
    return 0;
}