[HackerRank][Python3] Waiter

Problem :

https://www.hackerrank.com/challenges/waiter/problem

My Solution :

#!/usr/bin/env python3


def get_primes(q):
    primes = [2]
    i = 3
    while len(primes) < q:
        is_prime = True
        for prime in primes:
            if i % prime == 0:
                is_prime = False
                break
            if i < prime**2:
                break
        if is_prime:
            primes.append(i)
        i += 2
    return primes


def waiter(number, q):
    primes = get_primes(q)
    stack_a = number
    for i in range(q):
        stack_before = stack_a
        stack_a = []
        stack_b = []
        while stack_before:
            num = stack_before.pop()
            if num % primes[i] == 0:
                stack_b.append(num)
            else:
                stack_a.append(num)
        while stack_b:
            yield stack_b.pop()
    while stack_a:
        yield stack_a.pop()


n, q = map(int, input().split())
number = list(map(int, input().rstrip().split()))
result = waiter(number, q)
for num in result:
    print(num)

Comment :

stack 문제인지 소수 구하는 문제인지...