Description

Submission

• The Explanation of the second case is not correct

#include <bits/stdc++.h>

using namespace std;

const int prime_upper = 20000;
bool isPrime[prime_upper];

vector<int> calPrimes()
{
    vector<int> res;
    res.push_back(0); // 1-indexed result
    memset(isPrime, true, sizeof(isPrime));
    for(int i = 2; i < prime_upper; ++i) {
        if(isPrime[i]) {
            res.push_back(i);
            for(int j = 1; j * i < prime_upper; ++j) {
                isPrime[j * i] = false;
            }
        }
    }
    return res;
}


/*
 * Complete the waiter function below.
 */
vector<int> waiter(vector<int> number, int q) {
    /*
     * Write your code here.
     */
    vector<int> res;

    vector<int> primes = calPrimes();

    vector<stack<int>> B;
    stack<stack<int>> A;
    stack<int> A0;

    for(auto e : number) {
        A0.push(e);
    }

    A.push(A0);

    for(int i = 1; i <= q; ++i) {
        stack<int> Bi;
        stack<int> Ai;
        while (!A.top().empty()) {
            int e = A.top().top();
            A.top().pop();
            if (e % primes[i] == 0) {
                Bi.push(e);
            } else {
                Ai.push(e);
            }
        }
        B.push_back(Bi);
        A.push(Ai);
    }

    for(int i = 0; i < B.size(); ++i) {
        while(!B[i].empty()) {
            res.push_back(B[i].top());
            B[i].pop();
        }
    }

    while(!A.top().empty()) {
        res.push_back(A.top().top());
        A.top().pop();
    }

    return res;
}

int main()
{
    int n, q; cin >> n >> q;

    vector<int> number(n);

    for (int number_itr = 0; number_itr < n; number_itr++) {
        cin >> number[number_itr];
    }

    vector<int> result = waiter(number, q);

    for (int result_itr = 0; result_itr < result.size(); result_itr++) {
        cout << result[result_itr];

        if (result_itr != result.size() - 1) {
            cout << "\n";
        }
    }

    cout << "\n";

    return 0;
}