kd-treeの構築と近傍要素の抽出

// 木を分割する中央のインデクスを求める
int select_median_index(int n) {
  // 左詰にするため、あれこれ計算
  int h = log2(n);
  int s = 1 << h;
  int d = n - s;
  int s2 = s / 2;
  if (s2 > 0 && d >= s2) d = s2-1;
  return s2 + d;
}

// 「フォトンマッピング」本6.4 最近傍フォトンの効率的発見、locate_photons()参照
float traverse(PhotonAndDistance[] buf, vec pos, int index, float d2) {
  if (index < photons.length) {
    if (index * 2 + 2 < photons.length) {  // 子供がいる
      float e = axis_distance(nodes[index].divaxis, pos, photons[index].pos);
      if (e < 0) {
        d2 = traverse(buf, pos, index * 2 + 1, d2);  // 左の枝を見る
        if (e * e < d2) {
          d2 = traverse(buf, pos, index * 2 + 2, d2);  // 右の枝を見る
        }
      } else {
        d2 = traverse(buf, pos, index * 2 + 2, d2);  // 右の枝を見る
        if (e * e < d2) {
          d2 = traverse(buf, pos, index * 2 + 1, d2);  // 左の枝を見る
        }
      }
    }

    float e2 = photons[index].pos.sub(pos).sqlength();
    if (e2 < d2) {
      // フォトンを挿入
      int n = buf.length;
      int l, r;
      for (l=0, r=n; l<r; ) {
        int m = (l + r) / 2;
        if (e2 < buf[m].distance2) {
          r = m;
        } else {
          l = m + 1;
        }
      }
      // 後ろにずらす
      PhotonAndDistance temp = buf[buf.length-1];
      for (int i=buf.length; --i>l; ) buf[i] = buf[i-1];
      buf[l] = temp;
      buf[l].photon = photons[index];
      buf[l].distance2 = e2;
      // 最大距離更新
      d2 = buf[buf.length-1].distance2;
    }
  }
  return d2;
}