C++でパーティションを実装する - プログラミング原人の進化論

高速ソートアルゴリズムに関して、今回はパーティションを実装します。パーティションはクイックソートで利用されています。
流れとしては、パーティションの説明と実装の後、最後に問題を１題解いてみます。

目標

パーティションを実装できるようになる

パーティションとは

パーティションは、配列の末尾の要素を基準として、それ以下の大きさの要素とそれより大きい要素に分けるアルゴリズムです。末尾の要素をxとすると、x以下の要素を格納している領域とxより大きい要素を格納する領域を後方へ拡大していくイメージです。インデックスp~iは前者、インデックスi+1~j-1は後者の領域になっています。だいたい下のような感じです。

f:id:atuyusan:20190215164805p:plain

アルゴリズム

iをp-1、jをpにセットする。配列の末尾の要素をxとする。
以下jがr-1に達するまでを繰り返す：
1. jの位置にある要素がxより大きいならば、jを１大きくする。要素jがx以下ならば、iを１大きくしてから要素iと要素jを交換し、最後にjを１進める。
最後に、要素i+1と末尾を交換する。

要するに、配列の先頭から順番に大きさを検討して、x以下ならば「x以下のグループ」に、xより大きいならば「xより大きいグループ」に移動していくと言うことです。
計算量は $O(N)$ ですね。

実装してみた

与えられた配列Aに対してパーティション処理を行う関数を定義しています。なお、基準となる要素（もともと末尾であった要素）のインデックスを返します。

int partition(int A[], int p, int r) {
  int i, j, x, tmp;
  // initialize
  i = p - 1;
  x = A[r];
  
  for (j = p; j < r; j++) {
    if (A[j] <= x) {
      i++;
      // exchange
      tmp = A[i];
      A[i] = A[j];
      A[j] = tmp;
    }
  }
  // exchange
  tmp = A[i+1];
  A[i+1] = A[r];
  A[r] = tmp;

  return i+1;
}

問題を解く

上で定義した関数を用いて、Aizu Online Judgeの問題を解きます。
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_6_B

解答例：

#include<iostream>
#include<cstdio>
using namespace std;

int partition(int A[], int p, int r) {
  int i, j, x, tmp;
  // initialize
  i = p - 1;
  x = A[r];
  
  for (j = p; j < r; j++) {
    if (A[j] <= x) {
      i++;
      // exchange
      tmp = A[i];
      A[i] = A[j];
      A[j] = tmp;
    }
  }
  // exchange
  tmp = A[i+1];
  A[i+1] = A[r];
  A[r] = tmp;

  return i+1;
}

int main() {
  int i, n, A[100000], x_idx;

  cin >> n;
  for (i = 0; i < n; i++) scanf("%d", &A[i]);

  x_idx = partition(A, 0, n-1);

  for (i = 0; i < n; i++) {
    if (i > 0) printf(" ");
    if (i == x_idx) printf("[%d]", A[i]);
    else printf("%d", A[i]);
  }
  cout << endl;

  return 0;
}