kdtree2.cpp

kd tree java implementation 2

    for (int j=0; j<dim; j++) {
      dis += squared( the_data[i][j] - qv[j]);
    }
    e.dis = dis;
    e.idx = i;
    result.push_back(e);
  }
  sort(result.begin(), result.end() ); 

}


void kdtree2::n_nearest(vector<float>& qv, int nn, kdtree2_result_vector& result) {
  searchrecord sr(qv,*this,result);
  vector<float> vdiff(dim,0.0); 

  result.clear(); 

  sr.centeridx = -1;
  sr.correltime = 0;
  sr.nn = nn; 

  root->search(sr); 

  if (sort_results) sort(result.begin(), result.end());
  
}
// search for n nearest to a given query vector 'qv'.

  
void kdtree2::n_nearest_around_point(int idxin, int correltime, int nn,
				     kdtree2_result_vector& result) {
  vector<float> qv(dim);  //  query vector

  result.clear(); 

  for (int i=0; i<dim; i++) {
    qv[i] = the_data[idxin][i]; 
  }
  // copy the query vector.
  
  {
    searchrecord sr(qv, *this, result);
    // construct the search record.
    sr.centeridx = idxin;
    sr.correltime = correltime;
    sr.nn = nn; 
    root->search(sr); 
  }

  if (sort_results) sort(result.begin(), result.end());
    
}


void kdtree2::r_nearest(vector<float>& qv, float r2, kdtree2_result_vector& result) {
// search for all within a ball of a certain radius
  searchrecord sr(qv,*this,result);
  vector<float> vdiff(dim,0.0); 

  result.clear(); 

  sr.centeridx = -1;
  sr.correltime = 0;
  sr.nn = 0; 
  sr.ballsize = r2; 

  root->search(sr); 

  if (sort_results) sort(result.begin(), result.end());
  
}

int kdtree2::r_count(vector<float>& qv, float r2) {
// search for all within a ball of a certain radius
  {
    kdtree2_result_vector result; 
    searchrecord sr(qv,*this,result);

    sr.centeridx = -1;
    sr.correltime = 0;
    sr.nn = 0; 
    sr.ballsize = r2; 
    
    root->search(sr); 
    return(result.size());
  }

  
}

void kdtree2::r_nearest_around_point(int idxin, int correltime, float r2,
				     kdtree2_result_vector& result) {
  vector<float> qv(dim);  //  query vector

  result.clear(); 

  for (int i=0; i<dim; i++) {
    qv[i] = the_data[idxin][i]; 
  }
  // copy the query vector.
  
  {
    searchrecord sr(qv, *this, result);
    // construct the search record.
    sr.centeridx = idxin;
    sr.correltime = correltime;
    sr.ballsize = r2; 
    sr.nn = 0; 
    root->search(sr); 
  }

  if (sort_results) sort(result.begin(), result.end());
    
}


int kdtree2::r_count_around_point(int idxin, int correltime, float r2) 
{
  vector<float> qv(dim);  //  query vector


  for (int i=0; i<dim; i++) {
    qv[i] = the_data[idxin][i]; 
  }
  // copy the query vector.
  
  {
    kdtree2_result_vector result; 
    searchrecord sr(qv, *this, result);
    // construct the search record.
    sr.centeridx = idxin;
    sr.correltime = correltime;
    sr.ballsize = r2; 
    sr.nn = 0; 
    root->search(sr); 
    return(result.size());
  }

  
}


// 
//        KDTREE2_NODE implementation
//

// constructor
kdtree2_node::kdtree2_node(int dim) : box(dim) { 
  left = right = NULL; 
  //
  // all other construction is handled for real in the 
  // kdtree2 building operations.
  // 
}

// destructor
kdtree2_node::~kdtree2_node() {
  if (left != NULL) delete left; 
  if (right != NULL) delete right; 
  // maxbox and minbox 
  // will be automatically deleted in their own destructors. 
}


void kdtree2_node::search(searchrecord& sr) {
  // the core search routine.
  // This uses true distance to bounding box as the
  // criterion to search the secondary node. 
  //
  // This results in somewhat fewer searches of the secondary nodes
  // than 'search', which uses the vdiff vector,  but as this
  // takes more computational time, the overall performance may not
  // be improved in actual run time. 
  //

  if ( (left == NULL) && (right == NULL)) {
    // we are on a terminal node
    if (sr.nn == 0) {
      process_terminal_node_fixedball(sr);
    } else {
      process_terminal_node(sr);
    }
  } else {
    kdtree2_node *ncloser, *nfarther;

    float extra;
    float qval = sr.qv[cut_dim]; 
    // value of the wall boundary on the cut dimension. 
    if (qval < cut_val) {
      ncloser = left;
      nfarther = right;
      extra = cut_val_right-qval;
    } else {
      ncloser = right;
      nfarther = left;
      extra = qval-cut_val_left; 
    };

    if (ncloser != NULL) ncloser->search(sr);

    if ((nfarther != NULL) && (squared(extra) < sr.ballsize)) {
      // first cut
      if (nfarther->box_in_search_range(sr)) {
	nfarther->search(sr); 
      }      
    }
  }
}


inline float dis_from_bnd(float x, float amin, float amax) {
  if (x > amax) {
    return(x-amax); 
  } else if (x < amin)
    return (amin-x);
  else
    return 0.0;
  
}

inline bool kdtree2_node::box_in_search_range(searchrecord& sr) {
  //
  // does the bounding box, represented by minbox[*],maxbox[*]
  // have any point which is within 'sr.ballsize' to 'sr.qv'??
  //
 
  int dim = sr.dim;
  float dis2 =0.0; 
  float ballsize = sr.ballsize; 
  for (int i=0; i<dim;i++) {
    dis2 += squared(dis_from_bnd(sr.qv[i],box[i].lower,box[i].upper));
    if (dis2 > ballsize)
      return(false);
  }
  return(true);
}


void kdtree2_node::process_terminal_node(searchrecord& sr) {
  int centeridx  = sr.centeridx;
  int correltime = sr.correltime;
  unsigned int nn = sr.nn; 
  int dim        = sr.dim;
  float ballsize = sr.ballsize;
  //
  bool rearrange = sr.rearrange; 
  const kdtree2_array& data = *sr.data;

  const bool debug = false;

  if (debug) {
    printf("Processing terminal node %d, %d\n",l,u);
    cout << "Query vector = [";
    for (int i=0; i<dim; i++) cout << sr.qv[i] << ','; 
    cout << "]\n";
    cout << "nn = " << nn << '\n';
    check_query_in_bound(sr);
  }

  for (int i=l; i<=u;i++) {
    int indexofi;  // sr.ind[i]; 
    float dis;
    bool early_exit; 

    if (rearrange) {
      early_exit = false;
      dis = 0.0;
      for (int k=0; k<dim; k++) {
	dis += squared(data[i][k] - sr.qv[k]);
	if (dis > ballsize) {
	  early_exit=true; 
	  break;
	}
      }
      if(early_exit) continue; // next iteration of mainloop
      // why do we do things like this?  because if we take an early
      // exit (due to distance being too large) which is common, then
      // we need not read in the actual point index, thus saving main
      // memory bandwidth.  If the distance to point is less than the
      // ballsize, though, then we need the index.
      //
      indexofi = sr.ind[i];
    } else {
      // 
      // but if we are not using the rearranged data, then
      // we must always 
      indexofi = sr.ind[i];
      early_exit = false;
      dis = 0.0;
      for (int k=0; k<dim; k++) {
	dis += squared(data[indexofi][k] - sr.qv[k]);
	if (dis > ballsize) {
	  early_exit= true; 
	  break;
	}
      }
      if(early_exit) continue; // next iteration of mainloop
    } // end if rearrange. 
    
    if (centeridx > 0) {
      // we are doing decorrelation interval
      if (abs(indexofi-centeridx) < correltime) continue; // skip this point. 
    }

    // here the point must be added to the list.
    //
    // two choices for any point.  The list so far is either
    // undersized, or it is not.
    //
    if (sr.result.size() < nn) {
      kdtree2_result e;
      e.idx = indexofi;
      e.dis = dis;
      sr.result.push_element_and_heapify(e); 
      if (debug) cout << "unilaterally pushed dis=" << dis;
      if (sr.result.size() == nn) ballsize = sr.result.max_value();
      // Set the ball radius to the largest on the list (maximum priority).
      if (debug) {
	cout << " ballsize = " << ballsize << "\n"; 
	cout << "sr.result.size() = "  << sr.result.size() << '\n';
      }
    } else {
      //
      // if we get here then the current node, has a squared 
      // distance smaller
      // than the last on the list, and belongs on the list.
      // 
      kdtree2_result e;
      e.idx = indexofi;
      e.dis = dis;
      ballsize = sr.result.replace_maxpri_elt_return_new_maxpri(e); 
      if (debug) {
	cout << "Replaced maximum dis with dis=" << dis << 
	  " new ballsize =" << ballsize << '\n';
      }
    }
  } // main loop
  sr.ballsize = ballsize;
}

void kdtree2_node::process_terminal_node_fixedball(searchrecord& sr) {
  int centeridx  = sr.centeridx;
  int correltime = sr.correltime;
  int dim        = sr.dim;
  float ballsize = sr.ballsize;
  //
  bool rearrange = sr.rearrange; 
  const kdtree2_array& data = *sr.data;

  for (int i=l; i<=u;i++) {
    int indexofi = sr.ind[i]; 
    float dis;
    bool early_exit; 

    if (rearrange) {
      early_exit = false;
      dis = 0.0;
      for (int k=0; k<dim; k++) {
	dis += squared(data[i][k] - sr.qv[k]);
	if (dis > ballsize) {
	  early_exit=true; 
	  break;
	}
      }
      if(early_exit) continue; // next iteration of mainloop
      // why do we do things like this?  because if we take an early
      // exit (due to distance being too large) which is common, then
      // we need not read in the actual point index, thus saving main
      // memory bandwidth.  If the distance to point is less than the
      // ballsize, though, then we need the index.
      //
      indexofi = sr.ind[i];
    } else {
      // 
      // but if we are not using the rearranged data, then
      // we must always 
      indexofi = sr.ind[i];
      early_exit = false;
      dis = 0.0;
      for (int k=0; k<dim; k++) {
	dis += squared(data[indexofi][k] - sr.qv[k]);
	if (dis > ballsize) {
	  early_exit= true; 
	  break;
	}
      }
      if(early_exit) continue; // next iteration of mainloop
    } // end if rearrange. 
    
    if (centeridx > 0) {
      // we are doing decorrelation interval
      if (abs(indexofi-centeridx) < correltime) continue; // skip this point. 
    }

    {
      kdtree2_result e;
      e.idx = indexofi;
      e.dis = dis;
	sr.result.push_back(e);
    }

  }
}