Zig-zag Algothrim using C++ FA

zig_zag.hpp

#include <vector>

/**
 * @brief Zig-zag Algothrim
 * @param input: vector<T>
 * @return vector<T> (zig-zag vector)
 */
template<class T>
std::vector<T> zig_zag(std::vector<T> input) {
    
    // check vector size
    size_t len = static_cast<size_t>(sqrt(input.size()));
    assert(input.size() == len * len);

    std::vector<T> to_ret;
    size_t row_idx = 0, col_idx = 0;

    enum STATE {
        INIT,   RIGHT,      LEFT_DOWN,
        DOWN,   RIGHT_UP,   END
    };

    STATE next_step = STATE::INIT;
    while(next_step != STATE::END) {

        to_ret.push_back(input[row_idx*len + col_idx]);
        
        // init
        if (next_step == STATE::INIT) {
            next_step = STATE::RIGHT;
            col_idx += 1;
        }
        // end
        else if (row_idx == (len-1) && col_idx == (len-1)) {
            next_step = STATE::END;
        }
        
        // corner
        else if (next_step == STATE::RIGHT_UP && col_idx == len-1 && row_idx == 0) {
            next_step = STATE::DOWN;
            row_idx += 1;
        }
        // corner
        else if (next_step == STATE::LEFT_DOWN && row_idx == len-1 && col_idx == 0) {
            next_step = STATE::RIGHT;
            col_idx += 1;
        }
        
        // left half || right half
        else if (
            (next_step == STATE::RIGHT && row_idx == 0) ||
            (next_step == STATE::DOWN && col_idx == len-1)
         ) {
            next_step = STATE::LEFT_DOWN;
            col_idx -= 1;
            row_idx += 1;
        }
        // left half || right half
        else if (
            (next_step == STATE::DOWN && col_idx == 0) ||
            (next_step == STATE::RIGHT && row_idx == len-1)
        ) {
            next_step = STATE::RIGHT_UP;
            row_idx -= 1;
            col_idx += 1;
        }
        // left half || right
        else if (
            (next_step == STATE::LEFT_DOWN && col_idx == 0) ||
            (next_step == STATE::RIGHT_UP && col_idx == len-1)
        ) {
            next_step = STATE::DOWN;
            row_idx += 1;
        }
        // left half || right half
        else if (
            (next_step == STATE::LEFT_DOWN && row_idx == len-1) ||
            (next_step == STATE::RIGHT_UP && row_idx == 0)
         ) {
            next_step = STATE::RIGHT;
            col_idx += 1;
        }
        
        // normal
        else if (next_step == STATE::LEFT_DOWN) {
            col_idx -= 1;
            row_idx += 1;
        }
        // normal
        else if (next_step == STATE::RIGHT_UP) {
            row_idx -= 1;
            col_idx += 1;
        }
        // catch error
        else
            std::cout << "error!" << std::endl;
    }
    return to_ret;
}