MO's algorithm

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

 

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

 

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

연습 문제는 아래와 같다.

 

자연수로 이루어진 $$시퀀스가 있다. 이 때 다음 쿼리를 처리해야 한다.

1. $$가 주어졌을 때 $$에 존재하는 수들 중 서로 다른 수의 개수를 출력한다.

 

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

 

완전 탐색

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

// 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;
}

 

이전 쿼리를 재활용

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

 

쿼리간에 겹치는 구간이 있다

 

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

1. 회색 원소들 : $$과 $$ 어디에도 속하지 않는 원소들
2. 파란색 원소들 : $$과 $$ 두 곳에 모두 속하는 원소들
3. 주황색 원소들 : $$에만 속하는 원소들
4. 초록색 원소들 : $$에만 속하는 원소들

$$에서 $$로 넘어갈 때 1번에 속하는 원소들은 신경쓰지 않아도 됨은 자명하다. 그렇다면 쿼리가 변경될 때 2, 3, 4번에 속하는 원소들만 적절하게 처리해 준다면 이전 쿼리의 내용을 일부 재활용 할 수 있을 것이다. 2번 원소들은 양쪽에 모두 속하므로 변경을 할 필요가 없다. 3번 원소들은 쿼리가 변경될 때 제거가 될 것이고 4번 원소들은 쿼리가 변경될 때 추가가 될 것이다. 이 때 이러한 추가, 제거 연산을 편하게 하기 위해 $$배열을 도입한다. 아래 코드는 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;
  }
}

 

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

  // 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;
  }

 

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

 

쿼리를 정렬

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

  // 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의 정당성

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

 

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

 

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

 

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

 

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

 

전체 코드

 

#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]);
  }
}