遥かへのスピードランナー

シリコンバレーでAndroidアプリの開発してます。コンピュータービジョン・3D・アルゴリズム界隈にもたまに出現します。

heap sortを使った上位ランキング取得プログラム

Managing Gigabytes 4.6章で解説されているソートのプログラムを実装してみた。

検索エンジンなどでN個のデータの中から上位r個を取得したい場合、まずN個のデータからなるmax-heapを構成して、ルート(最大値)から順にr個をヒープから取り除くというアプローチが考えられる。しかしN>>rの場合、r個のデータからなるmin-heapを構成して、残りのN-r個のデータをheapのルート(最小値)と順次比較して、ルートよりも大きい場合はルートと入れ替えて、heapを再構成する、という手順を取った方がより計算量が少なくて済む、という話。

10万件のランダムな数値をスペース区切りで出力するプログラム
#include <iostream>
#include <stdlib.h>
using namespace std;

int main (int argc, char *argv[]) {
  srand((unsigned int)time(0));
  for (int i=0; i<100000; i++) {
    int tmp = rand()/1000;
    cout << tmp << " ";
  }
  cout << endl;
  return 1;
}
10件のmin-heapを構成して、10万件のデータから上位10件を取得・表示するプログラム
#include <iostream>
#include <fstream>
#define SWAP(a,b) (a ^= b,b = a ^ b,a ^= b)

using namespace std;

#define HEAPSIZE 10
#define DATAN 100000

int heapsize = HEAPSIZE;
int heap[HEAPSIZE];
int datan = DATAN;

void make_heap (int root);

int main (int argc, char *argv[]) {
  srand((unsigned int)time(0));
  ifstream input_file("testdata");
  int testdata[datan];
  for (int i=0; i<datan; i++) {
    input_file >> testdata[i];
  }

  //copy the first r accumlators into the heap
  for (int i=0; i<heapsize; i++) {
    heap[i] = testdata[i];
  }

  //convert array into min-heap
  for (int i=((heapsize+1)/2); i>0; i--) {
    make_heap(i-1);
  }

  for (int i=heapsize; i<datan; i++) {
    //testdata[i];
    if (heap[0] < testdata[i]) {
      //discard rheap[0] and insert testdata[i] into min-heap
      heap[0] = testdata[i];
      make_heap(0);
    }
  }

  //display top-r number in the ascending order
  while (heapsize > 0) {
    cout << heap[0] << endl;
    heap[0] = heap[heapsize-1];
    make_heap(0);
    heapsize--;
  }

  return 1;
}

void make_heap (int root) {
  int left = (root+1)*2-1;
  int right = left+1;
  if (left >= heapsize) {
    return;
  }
  if (heap[root] <= heap[left]) {
    if (right >= heapsize || heap[root] <= heap[right]) {
      //root is min
      //do nothing
    } else {
      //right is min
      //printf("swap %d and %d\n", root, right);
      SWAP(heap[root], heap[right]);
      make_heap(right);
    }
  } else {
    if (right >= heapsize || heap[left] <= heap[right]) {
      //left is min
      //printf("swap %d and %d\n", root, left);
      SWAP(heap[root], heap[left]);
      make_heap(left);
    } else {
      //right is min
      //printf("swap %d and %d\n", root, right);
      SWAP(heap[root], heap[right]);
      make_heap(right);
    }
  }
}

考察とかは特になしです。他のアルゴリズムとのベンチ比較とかも時間があればやってみたい。