Generalized posfixed quadratic equation solver

generalized-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,y) = (ux + p) * (vx + q) + r,
 * the algorithm is same as before, only the coefficients are computed in
 * a different way (see the function getProdFormCoeffs()).
 */

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 variables 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";
    // (ux + p) * (vx + q) + r form, becomes (ux + p) * (vy + q) + r
    else if (xpos.size() == 2) expr[xpos[1]] = "y";
    return expr;
}

// Function to detect if a postfix expression is in product form
bool isProductForm(vector<string>& expr) {
    // count x, count y, count sum
    int countx = 0, county = 0, counts = 0;
    for(size_t i = 0 ; i < expr.size() ; i++) {
        if(expr[i] == "x") countx++;
        if(expr[i] == "y") county++;
        if(expr[i] == "+" or expr[i] == "-") counts++;
    }
    return (countx == 1 and county == 1 and (counts == 2 or counts == 3));
}

// Find a, b, and c coefficients from f(x,y)
vector<float> getQuadFormCoeffs(vector<string>& expr) {
    /**
     * f(x,y) = ax^2 + by + c
     * f(0,0) = c
     * f(1,0) = a + c, therefore, a = f(1,0) - c
     * f(0,1) = b + c, therefore, b = f(0,1) - c
     */
    float c = eval(expr, 0.0, 0.0); // c = f(0,0)
    float a = eval(expr, 1.0, 0.0) - c; // a = f(1,0) - c
    float b = eval(expr, 0.0, 1.0) - c; // b = f(0,1) - c
    return {a, b, c};
}

// Find a, b, and c coefficients from f(x,y) product form
vector<float> getProdFormCoeffs(vector<string>& expr) {
    /**
     * f(x,y) = (ux + p) * (vy + q) + r
     * f(x,y) = uvxy + (uqx + vpy) + (pq + r)
     *
     * f(0,0) = pq + r, threfore, c = f(0,0)
     * f(1,0) = uq + (pq + r)
     * f(0,1) = vp + (pq + r), therefore, b = f(1,0) + f(0,1) - 2c = uq + vp
     * f(1,1) = uv + (uq + vp) + (pq + r), therfore, a = f(1,1) - b - c
     */
    float c = eval(expr, 0.0, 0.0); // c = f(0,0)
    float b = eval(expr, 1.0, 0.0) + eval(expr, 0.0, 1.0) - (2.0 * c); // b = f(1,0) + f(0,1) - 2c
    float a = eval(expr, 1.0, 1.0) - b - c; // a = f(1,1) - b - c
    return {a, b, c};
}

// 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:
        // (ux + p) * (vx + q) + r
        {"2", "x", "*", "3", "+", "3", "x", "*", "2", "+", "*", "9", "-"},
        // (x + p) * (vx + q) + r
        {"x", "7", "-", "3", "x", "*", "2", "-", "*", "9", "+"},
        // (ux + p) * (x + q) + r
        {"5", "x", "*", "3", "-", "x", "3", "+", "*", "-9", "+"},
        // (x + p) * (x + q) + r
        {"x", "3", "+", "x", "2", "+", "*", "9", "-"},
        // (x + p) * (x + q)
        {"x", "3", "+", "x", "2", "+", "*"},
    };

    for(auto fx : quadratics) {
        vector<float> coeffs;
        cout << "f(x) = " << fx << endl;

        vector<string> fxy = renameVariables(fx);
        cout << "f(x,y) = " << fxy << ", product form = " << isProductForm(fxy) << endl;

        if(!isProductForm(fx))
            coeffs = getQuadFormCoeffs(fxy);
        else
            coeffs = getProdFormCoeffs(fxy);

        cout << "{a, b, c} = " << coeffs << endl;

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