MO's algorithm

MO's algorithm은 온라인으로 풀기 힘든 쿼리 문제를 오프라인으로 쉽게 풀 수 있게 해주는 알고리즘이다. MO's algorithm은 쿼리를 정렬한 뒤 정렬된 순서대로 처리를 한다. 이론은 이게 끝이다.

이론만 들어서는 아마 문제에 어떻게 적용해야 할 지 감이 안올것이다. 그럼 연습문제를 풀면서 MO's algorithm을 익혀보도록 하자. 연습문제는 https://www.acmicpc.net/problem/13547에서 풀 수 있다.

구간에 존재하는 서로 다른 수의 개수

---

연습 문제는 아래와 같다.

자연수로 이루어진 \(A_1,\ A_2,\ A_3,\ ...\ ,\ A_{N-1},\ A_{N}\)시퀀스가 있다. 이 때 다음 쿼리를 처리해야 한다.

1. \(x,\ y\ (1 \le x \le y \le N)\)가 주어졌을 때 \([A_x,\ A_y]\)에 존재하는 수들 중 서로 다른 수의 개수를 출력한다.

문제 자체는 상당히 단순하다. 문제를 꼬아놓은 것도 없으며 요구하는것도 명확하다. 이 문제에 어떻게 MO's algorithm을 적용할 수 있을까?

완전 탐색

---

우선 가장 쉽게 짤 수 있는 완전 탐색코드를 살펴 보자. 쿼리로 들어온 구간을 모두 살펴보면서 flag배열을 하나 둬 중복되는 것을 제외하고 수를 카운팅한다. 이 코드의 시간복잡도는 어떻게 될까? 쿼리의 개수가 \(Q\)개라 한다면 총 \(O(QN)\)의 시간복잡도를 가진다. \(Q\)나 \(N\)이 크다면 시간초과를 받기 딱 좋을것이다.

// brute force function
int query(int lo, int hi) {
  int ret = 0;
  
  // A[i]가 체크가 안됐다면 체크를 해주고 count를 1증가 시킨다
  for (int i = lo ; i <= hi ; ++i) {
    if (!isExist[A[i]]) {
      ++ret;
      isExist[A[i]] = true;
    }
  }
  
  // 다음 쿼리를 위해 true로 체크되어 있는 곳을 false로 초기화 시켜준다
  for (int i = lo ; i <= hi ; ++i) {
    isExist[A[i]] = false;
  }
  
  return ret;
}

이전 쿼리를 재활용

---

우리는 완전 탐색보다 조금 더 똑똑하게 코드를 짤 수 있다. 잘 생각해 본다면 이전 쿼리의 내용을 일부 재활용 할 수 있다는 것을 알 수 있다. 만약 \(Q_{i - 1}\)과 \(Q_i\)의 구간에 겹치는 부분이 있다고 생각을 해 보자. 그림으로 표현하자면 아래와 같다.

시퀀스의 원소들을 총 네가지로 분류할 수 있다.

1. 회색 원소들 : \(Q_{i - 1}\)과 \(Q_i\) 어디에도 속하지 않는 원소들

2. 파란색 원소들 : \(Q_{i - 1}\)과 \(Q_i\) 두 곳에 모두 속하는 원소들

3. 주황색 원소들 : \(Q_{i - 1}\)에만 속하는 원소들

4. 초록색 원소들 : \(Q_i\)에만 속하는 원소들

\(Q_{i - 1}\)에서 \(Q_i\)로 넘어갈 때 1번에 속하는 원소들은 신경쓰지 않아도 됨은 자명하다. 그렇다면 쿼리가 변경될 때 2, 3, 4번에 속하는 원소들만 적절하게 처리해 준다면 이전 쿼리의 내용을 일부 재활용 할 수 있을 것이다. 2번 원소들은 양쪽에 모두 속하므로 변경을 할 필요가 없다. 3번 원소들은 쿼리가 변경될 때 제거가 될 것이고 4번 원소들은 쿼리가 변경될 때 추가가 될 것이다. 이 때 이러한 추가, 제거 연산을 편하게 하기 위해 \(cnt\)배열을 도입한다. 아래 코드는 3번과 4번 연산을 어떻게 구현하는지 보여준다.

// add function
void add(int num) {
  // num이 새롭게 cnt배열에 들어올때만 카운팅 해준다
  if (++cnt[num] == 1) {
    ++distinctNumbers;
  }
}

// erase function
void erase(int num) {
  // num이 cnt배열에서 완전히 없어질때 카운팅 해준다
  if (--cnt[num] == 0) {
    --distinctNumbers;
  }
}

추가 함수에서는 \(cnt\)배열에 수가 처음 들어갈 때만 카운팅 해주고 제거 함수에서는 \(cnt\)배열에서 수가 완전히 없어졌을 때만 카운팅 해준다. 두 함수를 통해 이전 쿼리에서 다음 쿼리로 넘어갈 때 겹치는 부분을 재활용 할 수 있게 됐다. 재활용을 하는 코드는 어떻게 작성할 수 있을까? 이전 쿼리의 범위에서 다음 쿼리의 범위로 넘어갈 때 이터레이터를 조정해 가면서 적절하게 \(add\)함수와 \(erase\)함수를 호출해 주면 된다. 이 부분이 살짝 까다로울 수 있다. 이 글에서 쿼리의 범위를 나타내는 이터레이터는 \([lo,\ hi]\)로 표현된다.(양쪽이 전부 폐구간임을 주의해야 한다.)

  // Change the iterator from the previous query to the current query
  int lo = 0, hi = 0;
  
  // 시작 이터레이터는 [0, 0]이기 때문에 A[0]의 값을 cnt배열에 넣고 시작한다
  add(A[0]);
  
  // 모든 쿼리에 대한 답을 계산한다
  for (auto cur : query) {
    // 현재 위치는 [lo, hi]이고 다음 위치는 [nlo, nhi]이다
    int nlo = cur.lo, nhi = cur.hi;
    
    // lo를 nlo로 움직이는 부분
    // 1. lo가 nlo보다 왼쪽에 있는 경우 : 수를 제거할 필요가 있다
    for (int i = lo ; i < nlo ; ++i) {
      erase(A[i]);
    }

    // 2. lo가 nlo보다 오른쪽에 있는 경우 : 수를 더할 필요가 있다
    for (int i = lo - 1; i >= nlo ; --i) {
      add(A[i]);
    }
    
    // hi를 nhi로 움직이는 부분
    // 1. hi가 nhi보다 왼쪽에 있는 경우 : 수를 더할 필요가 있다
    for (int i = hi + 1 ; i <= nhi ; ++i) {
      add(A[i]);
    }

    // 2. hi가 nhi보다 오른쪽에 있는 경우 : 수를 제거할 필요가 있다
    for (int i = hi ; i > nhi ; --i) {
      erase(A[i]);
    }
    
    // 이터레이터가 현재 쿼리의 범위를 나타내도록 한다
    lo = nlo;
    hi = nhi;
    
    // 답을 저장할 배열에 실제 쿼리의 순서에 맞게 답을 기록해 준다
    ans[cur.idx] = distinctNumbers;
  }

이제 완벽하게 이전 쿼리를 재활용 할 수 있게 됐다. 그렇다면 이 코드의 시간복잡도는 어떻게 될까? 너무나 불행하게도 이코드 또한 \(O(QN)\)이다. 아주 최악의 경우에는 쿼리가 시퀀의 처음부분에 등장하고 끝 부분에 등장하고 다시 처음 부분에 등장하는 형태로 들어올 수 있기 때문이다. 그렇다면 우리는 이제까지 쓸모없는 코드를 짠 것일까?

쿼리를 정렬

---

놀랍게도 쿼리를 정렬함으로써 위 코드의 시간복잡도를 개선할 수 있다. 쿼리를 정렬할 때 Sqrt decomposition의 개념이 들어간다. 우선 시퀀스의 인덱스들을 일정한 크기로 나눈 뒤 나눠진 부분을 한 묶음으로 생각한다. 통상적으로 나누는 크기는 \(\sqrt N\)을 사용한다. 쿼리의 끝 부분이 어떠한 묶음에 놓여져 있는지 계산한 뒤 쿼리를 쿼리가 속한 묶음의 번호 순으로 정렬을 해 준다. 이 부분은 글을 읽는것보다 코드로 보는것이 훨씬 명쾌할 것이다.

  // sort queries
  int sz = sqrt(n);
  sort(query.begin(), query.end(), [&](const Query& l, const Query& r) {
    // 쿼리가 끝나는 부분이 어떠한 묶음에 속하는지 확인하다
    // 묶음의 인덱스가 더 작은곳에 속하는 쿼리를 먼저 처리한다
    // 만약 묶음의 인덱스가 같다면 쿼리의 시작지점이 작은것을 먼저 처리한다
    int lIdx = l.hi / sz;
    int rIdx = r.hi / sz;
    return lIdx == rIdx ? l.lo < r.lo : lIdx < rIdx;
  });

MO's algorithm의 정당성

---

이렇게 쿼리를 정렬하는 것만으로 \(O(QN)\)의 가망없는 시간복잡도가 개선될 수 있을까? 만약 된다면 왜 이러한 일이 벌어지는 것일까? 위에 코드를 한번 살펴보자. 정렬된 쿼리를 순회하면서 프로그램은 네개의 반복문을 실행한다. 처음 두개는 이터레이터 \(lo\)를 변경해주는 반복문이였고 그 다음 두개는 이터레이터 \(hi\)를 변경해주는 반복문이였다. 프로그램의 전체 수행시간은 모든 쿼리를 순회하면서 이 네 개의 반복문이 얼마나 실행이 됐는지에 따라 달라질 것이다.

먼저 이터레이터 \(lo\)가 얼마나 움직일 지 생각을 해보자. 같은 묶음 내에서 \(lo\)는 오름차순으로 정렬이 되어 있다. 이 때 \(lo\)는 \(A_1\)부터 \(A_N\)까지 움직일 수 있다. 즉 한 묶음 내에서 최대 \(O(N)\)만큼 움직일 수 있다는 얘기다. 우리는 한 묶음의 크기를 \(\sqrt N\)개로 정의 했으니 전체 묶음의 개수는 \(\sqrt N\)개가 될 것이고 \(lo\)는 최악의 경우 \(O(N \sqrt N)\)만큼 움직일 것이다.

이제 이터레이터 \(hi\)를 생각해보자. 이터레이터 \(hi\)는 묶음의 인덱스 순으로 정렬이 되어있다. 한 쿼리에서 다음 쿼리로 갈 때 최악의 경우 묶음의 크기인 \(\sqrt N\)만큼 움직일 수 있다. 쿼리의 개수가 \(Q\)라고 했을 때 \(O(Q \sqrt N)\)만큼 움직일 것이다.

이터레이터 \(lo\)와 \(hi\)의 시간복잡도를 계산했으니 이 두개의 시간복잡도를 합쳐주면 전체 프로그램의 수행시간이 된다. 즉 MO's algorithm을 사용하면 \(O((N\ + \ Q) \sqrt N)\)으로 이 문제를 해결할 수 있다. \(logN\)과 같이 기가막히게 아름다운 시간복잡도는 아니지만 \(O(QN)\)이 걸리던 기존의 알고리즘에 비하면 충분히 아름다운 시간복잡도라 할 수 있다.

마지막으로 이 문제의 전체 소스코드를 첨부한다.

전체 코드

---

#include <stdio.h>
#include <math.h>
#include <vector>
#include <algorithm>

using namespace std;

struct Query {
  int lo, hi, idx;
  Query(int lo, int hi, int idx) : lo(lo), hi(hi), idx(idx) {}
};

int n, m, A[100010], cnt[1000010], ans[100010], distinctNumbers;
vector<Query> query;

// add function
void add(int num) {
  // num이 새롭게 cnt배열에 들어올때만 카운팅 해준다
  if (++cnt[num] == 1) {
    ++distinctNumbers;
  }
}

// erase function
void erase(int num) {
  // num이 cnt배열에서 완전히 없어질때 카운팅 해준다
  if (--cnt[num] == 0) {
    --distinctNumbers;
  }
}

int main() {
  scanf("%d", &n);
  for (int i = 0 ; i < n ; ++i) {
    scanf("%d", &A[i]);
  }
  scanf("%d", &m);
  for (int i = 0 ; i < m ; ++i) {
    int lo, hi;
    scanf("%d%d", &lo, &hi);
    query.emplace_back(lo - 1, hi - 1, i);
  }
  
  // sort queries
  int sz = sqrt(n);
  sort(query.begin(), query.end(), [&](const Query& l, const Query& r) {
    // 쿼리가 끝나는 부분이 어떠한 묶음에 속하는지 확인하다
    // 묶음의 인덱스가 더 작은곳에 속하는 쿼리를 먼저 처리한다
    // 만약 묶음의 인덱스가 같다면 쿼리의 시작지점이 작은것을 먼저 처리한다
    int lIdx = l.hi / sz;
    int rIdx = r.hi / sz;
    return lIdx == rIdx ? l.lo < r.lo : lIdx < rIdx;
  });
  
  // Change the iterator from the previous query to the current query
  int lo = 0, hi = 0;
  
  // 시작 이터레이터는 [0, 0]이기 때문에 A[0]의 값을 cnt배열에 넣고 시작한다
  add(A[0]);
  
  // 모든 쿼리에 대한 답을 계산한다
  for (auto cur : query) {
    // 현재 위치는 [lo, hi]이고 다음 위치는 [nlo, nhi]이다
    int nlo = cur.lo, nhi = cur.hi;
    
    // lo를 nlo로 움직이는 부분
    // 1. lo가 nlo보다 왼쪽에 있는 경우 : 수를 제거할 필요가 있다
    for (int i = lo ; i < nlo ; ++i) {
      erase(A[i]);
    }
    
    // 2. lo가 nlo보다 오른쪽에 있는 경우 : 수를 더할 필요가 있다
    for (int i = lo - 1; i >= nlo ; --i) {
      add(A[i]);
    }
    
    // hi를 nhi로 움직이는 부분
    // 1. hi가 nhi보다 왼쪽에 있는 경우 : 수를 더할 필요가 있다
    for (int i = hi + 1 ; i <= nhi ; ++i) {
      add(A[i]);
    }
    
    // 2. hi가 nhi보다 오른쪽에 있는 경우 : 수를 제거할 필요가 있다
    for (int i = hi ; i > nhi ; --i) {
      erase(A[i]);
    }
    
    // 이터레이터가 현재 쿼리의 범위를 나타내도록 한다
    lo = nlo;
    hi = nhi;
    
    // 답을 저장할 배열에 실제 쿼리의 순서에 맞게 답을 기록해 준다
    ans[cur.idx] = distinctNumbers;
  }
  
  for (int i = 0 ; i < m ; ++i) {
    printf("%d\n", ans[i]);
  }
}