C++로 구현하는 지그재그 반복자(Zigzag Iterator)

C++ 지그재그 반복자(Zigzag Iterator)란?

두 개의 1차원 배열(벡터)이 주어졌을 때, 두 배열의 원소를 번갈아가며 차례로 반환하는 반복자를 구현하는 문제입니다. 이 반복자는 다음 두 가지 메서드를 제공해야 합니다.

• next() — 다음 원소를 반환합니다.
• hasNext() — 아직 반환할 다음 원소가 남아 있는지 확인합니다.

예를 들어 입력이 v1 = [1, 2], v2 = [3, 4, 5, 6]이라면 출력 순서는 [1, 3, 2, 4, 5, 6]처럼 두 배열을 교차하며 순회하게 됩니다. 한쪽 배열이 먼저 소진되더라도 나머지 배열의 원소는 끝까지 모두 출력되어야 합니다.

문제 해결 접근 방법

이 문제는 큐(queue)를 이용하면 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 각 배열마다 '현재 읽고 있는 위치(인덱스)'를 pair 형태로 관리하고, 이를 큐에 넣어 번갈아 처리하는 것입니다.

알고리즘 단계는 다음과 같습니다.

1. pair<int, int> 타입의 큐 q를 선언합니다. first는 배열 내 인덱스, second는 어느 배열인지 식별하는 값(0 또는 1)입니다.
2. 생성자에서 두 배열 v1, v2를 전달받습니다.
3. v1의 크기가 0보다 크면 {0, 0}을 큐에 삽입합니다.
4. v2의 크기가 0보다 크면 {0, 1}을 큐에 삽입합니다.
5. next() 함수에서는 큐 맨 앞의 원소를 꺼내(temp) 처리합니다.
6. temp.second가 1이라면(v2 배열):
   • ret := v2[temp.first]
   • temp.first를 1 증가시킵니다.
   • 증가된 인덱스가 아직 v2의 크기 미만이라면 temp를 다시 큐에 삽입합니다.
7. 그렇지 않다면(v1 배열):
   • ret := v1[temp.first]
   • temp.first를 1 증가시킵니다.
   • 증가된 인덱스가 아직 v1의 크기 미만이라면 temp를 다시 큐에 삽입합니다.
8. ret 값을 반환합니다.
9. hasNext() 함수는 큐가 비어 있지 않으면 true를 반환합니다.

이 방식 덕분에 한쪽 배열이 먼저 끝나면 해당 pair는 더 이상 큐에 들어가지 않으므로, 남은 배열의 원소만 자연스럽게 계속 출력됩니다.

예제 코드

다음은 위 알고리즘을 C++로 구현한 코드입니다.

#include <bits/stdc++.h>
using namespace std;

class ZigzagIterator {
public:
    queue<pair<int, int>> q;
    vector<int> v1, v2;

    ZigzagIterator(vector<int>& v1, vector<int>& v2) {
        this->v1 = v1;
        this->v2 = v2;
        if (v1.size()) {
            q.push({0, 0});
        }
        if (v2.size()) {
            q.push({0, 1});
        }
    }

    int next() {
        pair<int, int> temp;
        temp = q.front();
        q.pop();
        int ret = 0;
        if (temp.second == 1) {
            ret = v2[temp.first];
            temp.first++;
            if (temp.first < v2.size())
                q.push(temp);
        } else {
            ret = v1[temp.first];
            temp.first++;
            if (temp.first < v1.size())
                q.push(temp);
        }
        return ret;
    }

    bool hasNext() {
        return !q.empty();
    }
};

main(){
    vector<int> v1 = {1,3,5,7}, v2 = {2,4,6,8,10,12,17};
    ZigzagIterator ob(v1, v2);
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.hasNext() ? "True" : "False") << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.hasNext() ? "True" : "False") << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.hasNext() ? "True" : "False") << endl;
    cout << (ob.next()) << endl;
    cout << (ob.next()) << endl;
    cout << (ob.hasNext() ? "True" : "False") << endl;
}

입력

{1,3,5,7},{2,4,6,8,10,12,17}

출력

1
2
True
3
4
5
True
6
7
8
10
True
12
17
False

동작 원리 살펴보기

v1 = {1, 3, 5, 7}, v2 = {2, 4, 6, 8, 10, 12, 17}일 때의 동작 과정을 정리하면 다음과 같습니다.

1. 생성자에서 큐에는 {0, 0}(v1용)과 {0, 1}(v2용)이 순서대로 들어갑니다.
2. next() 호출 시 큐의 앞에서 원소를 하나 꺼내 값을 반환하고, 아직 남은 원소가 있다면 인덱스를 1 늘려 다시 큐 뒤에 넣습니다.
3. 이 과정이 반복되면서 1 → 2 → 3 → 4 → 5 → 6 → 7 → 8 순으로 두 배열이 번갈아 출력됩니다.
4. v1이 먼저 소진되면(원소 4개) 이후부터는 v2의 나머지 원소인 10, 12, 17만 연속해서 출력됩니다.
5. 모든 원소를 소비하면 hasNext()가 false를 반환하며 순회가 종료됩니다.

마치며

이 구현은 next() 호출당 O(1)의 시간 복잡도를 가지며, 큐에는 배열 수만큼의 pair(여기서는 최대 2개)만 유지되므로 공간 복잡도 역시 O(1)입니다. 특히 큐 기반 구조는 세 개 이상의 배열로도 손쉽게 확장할 수 있다는 장점이 있어, 실전 코딩 테스트나 LeetCode 스타일의 문제에서 유용하게 활용됩니다.