kdtree.cpp

Go to the documentation of this file.

/******************************************************************************
 **  Filename:  kdtree.cpp
 **  Purpose:   Routines for managing K-D search trees
 **  Author:    Dan Johnson
 **
 **  (c) Copyright Hewlett-Packard Company, 1988.
 ** Licensed under the Apache License, Version 2.0 (the "License");
 ** you may not use this file except in compliance with the License.
 ** You may obtain a copy of the License at
 ** http://www.apache.org/licenses/LICENSE-2.0
 ** Unless required by applicable law or agreed to in writing, software
 ** distributed under the License is distributed on an "AS IS" BASIS,
 ** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 ** See the License for the specific language governing permissions and
 ** limitations under the License.
 ******************************************************************************/
 
/*-----------------------------------------------------------------------------
          Include Files and Type Defines
-----------------------------------------------------------------------------*/
#include "kdtree.h"
 
#include <algorithm>
#include <cfloat> // for FLT_MAX
#include <cmath>
#include <cstdio>
 
namespace tesseract {
 
#define Magnitude(X) ((X) < 0 ? -(X) : (X))
#define NodeFound(N, K, D) (((N)->Key == (K)) && ((N)->Data == (D)))
 
/*-----------------------------------------------------------------------------
        Global Data Definitions and Declarations
-----------------------------------------------------------------------------*/
#define MINSEARCH (-FLT_MAX)
#define MAXSEARCH FLT_MAX
 
// Helper function to find the next essential dimension in a cycle.
static int NextLevel(KDTREE *tree, int level) {
  do {
    ++level;
    if (level >= tree->KeySize) {
      level = 0;
    }
  } while (tree->KeyDesc[level].NonEssential);
  return level;
}
 
//-----------------------------------------------------------------------------
template <typename Key, typename Value>
class MinK {
public:
  MinK(Key max_key, int k);
  ~MinK();
 
  struct Element {
    Element() = default;
    Element(const Key &k, const Value &v) : key(k), value(v) {}
 
    Key key;
    Value value;
  };
 
  bool insert(Key k, Value v);
  const Key &max_insertable_key();
 
  int elements_count() {
    return elements_count_;
  }
  const Element *elements() {
    return elements_;
  }
 
private:
  const Key max_key_;  
  Element *elements_;  
  int elements_count_; 
  int k_;              
  int max_index_;      
};
 
template <typename Key, typename Value>
MinK<Key, Value>::MinK(Key max_key, int k)
    : max_key_(max_key), elements_count_(0), k_(k < 1 ? 1 : k), max_index_(0) {
  elements_ = new Element[k_];
}
 
template <typename Key, typename Value>
MinK<Key, Value>::~MinK() {
  delete[] elements_;
}
 
template <typename Key, typename Value>
const Key &MinK<Key, Value>::max_insertable_key() {
  if (elements_count_ < k_) {
    return max_key_;
  }
  return elements_[max_index_].key;
}
 
template <typename Key, typename Value>
bool MinK<Key, Value>::insert(Key key, Value value) {
  if (elements_count_ < k_) {
    elements_[elements_count_++] = Element(key, value);
    if (key > elements_[max_index_].key) {
      max_index_ = elements_count_ - 1;
    }
    return true;
  } else if (key < elements_[max_index_].key) {
    // evict the largest element.
    elements_[max_index_] = Element(key, value);
    // recompute max_index_
    for (int i = 0; i < elements_count_; i++) {
      if (elements_[i].key > elements_[max_index_].key) {
        max_index_ = i;
      }
    }
    return true;
  }
  return false;
}
 
//-----------------------------------------------------------------------------
class KDTreeSearch {
public:
  KDTreeSearch(KDTREE *tree, float *query_point, int k_closest);
  ~KDTreeSearch();
 
  void Search(int *result_count, float *distances, void **results);
 
private:
  void SearchRec(int Level, KDNODE *SubTree);
  bool BoxIntersectsSearch(float *lower, float *upper);
 
  KDTREE *tree_;
  float *query_point_;
  float *sb_min_; 
  float *sb_max_; 
  MinK<float, void *> results_;
};
 
KDTreeSearch::KDTreeSearch(KDTREE *tree, float *query_point, int k_closest)
    : tree_(tree), query_point_(query_point), results_(MAXSEARCH, k_closest) {
  sb_min_ = new float[tree->KeySize];
  sb_max_ = new float[tree->KeySize];
}
 
KDTreeSearch::~KDTreeSearch() {
  delete[] sb_min_;
  delete[] sb_max_;
}
 
void KDTreeSearch::Search(int *result_count, float *distances, void **results) {
  if (tree_->Root.Left == nullptr) {
    *result_count = 0;
  } else {
    for (int i = 0; i < tree_->KeySize; i++) {
      sb_min_[i] = tree_->KeyDesc[i].Min;
      sb_max_[i] = tree_->KeyDesc[i].Max;
    }
    SearchRec(0, tree_->Root.Left);
    int count = results_.elements_count();
    *result_count = count;
    for (int j = 0; j < count; j++) {
      // Pre-cast to float64 as key is a template type and we have no control
      // over its actual type.
      distances[j] = static_cast<float>(sqrt(static_cast<double>(results_.elements()[j].key)));
      results[j] = results_.elements()[j].value;
    }
  }
}
 
/*-----------------------------------------------------------------------------
              Public Code
-----------------------------------------------------------------------------*/
KDTREE *MakeKDTree(int16_t KeySize, const PARAM_DESC KeyDesc[]) {
  auto *KDTree = new KDTREE(KeySize);
  for (int i = 0; i < KeySize; i++) {
    KDTree->KeyDesc[i].NonEssential = KeyDesc[i].NonEssential;
    KDTree->KeyDesc[i].Circular = KeyDesc[i].Circular;
    if (KeyDesc[i].Circular) {
      KDTree->KeyDesc[i].Min = KeyDesc[i].Min;
      KDTree->KeyDesc[i].Max = KeyDesc[i].Max;
      KDTree->KeyDesc[i].Range = KeyDesc[i].Max - KeyDesc[i].Min;
      KDTree->KeyDesc[i].HalfRange = KDTree->KeyDesc[i].Range / 2;
      KDTree->KeyDesc[i].MidRange = (KeyDesc[i].Max + KeyDesc[i].Min) / 2;
    } else {
      KDTree->KeyDesc[i].Min = MINSEARCH;
      KDTree->KeyDesc[i].Max = MAXSEARCH;
    }
  }
  KDTree->Root.Left = nullptr;
  KDTree->Root.Right = nullptr;
  return KDTree;
}
 
void KDStore(KDTREE *Tree, float *Key, CLUSTER *Data) {
  auto PtrToNode = &(Tree->Root.Left);
  auto Node = *PtrToNode;
  auto Level = NextLevel(Tree, -1);
  while (Node != nullptr) {
    if (Key[Level] < Node->BranchPoint) {
      PtrToNode = &(Node->Left);
      if (Key[Level] > Node->LeftBranch) {
        Node->LeftBranch = Key[Level];
      }
    } else {
      PtrToNode = &(Node->Right);
      if (Key[Level] < Node->RightBranch) {
        Node->RightBranch = Key[Level];
      }
    }
    Level = NextLevel(Tree, Level);
    Node = *PtrToNode;
  }
 
  *PtrToNode = new KDNODE(Tree, Key, Data, Level);
} /* KDStore */
 
void KDDelete(KDTREE *Tree, float Key[], void *Data) {
  int Level;
  KDNODE *Current;
  KDNODE *Father;
 
  /* initialize search at root of tree */
  Father = &(Tree->Root);
  Current = Father->Left;
  Level = NextLevel(Tree, -1);
 
  /* search tree for node to be deleted */
  while ((Current != nullptr) && (!NodeFound(Current, Key, Data))) {
    Father = Current;
    if (Key[Level] < Current->BranchPoint) {
      Current = Current->Left;
    } else {
      Current = Current->Right;
    }
 
    Level = NextLevel(Tree, Level);
  }
 
  if (Current != nullptr) { /* if node to be deleted was found */
    if (Current == Father->Left) {
      Father->Left = nullptr;
      Father->LeftBranch = Tree->KeyDesc[Level].Min;
    } else {
      Father->Right = nullptr;
      Father->RightBranch = Tree->KeyDesc[Level].Max;
    }
 
    InsertNodes(Tree, Current->Left);
    InsertNodes(Tree, Current->Right);
    delete Current;
  }
} /* KDDelete */
 
void KDNearestNeighborSearch(KDTREE *Tree, float Query[], int QuerySize, float MaxDistance,
                             int *NumberOfResults, void **NBuffer, float DBuffer[]) {
  KDTreeSearch search(Tree, Query, QuerySize);
  search.Search(NumberOfResults, DBuffer, NBuffer);
}
 
/*---------------------------------------------------------------------------*/
void KDWalk(KDTREE *Tree, kdwalk_proc action, ClusteringContext *context) {
  if (Tree->Root.Left != nullptr) {
    Walk(Tree, action, context, Tree->Root.Left, NextLevel(Tree, -1));
  }
}
 
/*-----------------------------------------------------------------------------
              Private Code
-----------------------------------------------------------------------------*/
 
/*---------------------------------------------------------------------------*/
void KDTreeSearch::SearchRec(int level, KDNODE *sub_tree) {
  if (level >= tree_->KeySize) {
    level = 0;
  }
 
  if (!BoxIntersectsSearch(sb_min_, sb_max_)) {
    return;
  }
 
  results_.insert(DistanceSquared(tree_->KeySize, &tree_->KeyDesc[0], query_point_, sub_tree->Key),
                  sub_tree->Data);
 
  if (query_point_[level] < sub_tree->BranchPoint) {
    if (sub_tree->Left != nullptr) {
      float tmp = sb_max_[level];
      sb_max_[level] = sub_tree->LeftBranch;
      SearchRec(NextLevel(tree_, level), sub_tree->Left);
      sb_max_[level] = tmp;
    }
    if (sub_tree->Right != nullptr) {
      float tmp = sb_min_[level];
      sb_min_[level] = sub_tree->RightBranch;
      SearchRec(NextLevel(tree_, level), sub_tree->Right);
      sb_min_[level] = tmp;
    }
  } else {
    if (sub_tree->Right != nullptr) {
      float tmp = sb_min_[level];
      sb_min_[level] = sub_tree->RightBranch;
      SearchRec(NextLevel(tree_, level), sub_tree->Right);
      sb_min_[level] = tmp;
    }
    if (sub_tree->Left != nullptr) {
      float tmp = sb_max_[level];
      sb_max_[level] = sub_tree->LeftBranch;
      SearchRec(NextLevel(tree_, level), sub_tree->Left);
      sb_max_[level] = tmp;
    }
  }
}
 
/*---------------------------------------------------------------------------*/
float DistanceSquared(int k, PARAM_DESC *dim, float p1[], float p2[]) {
  float total_distance = 0;
 
  for (; k > 0; k--, p1++, p2++, dim++) {
    if (dim->NonEssential) {
      continue;
    }
 
    float dimension_distance = *p1 - *p2;
 
    /* if this dimension is circular - check wraparound distance */
    if (dim->Circular) {
      dimension_distance = Magnitude(dimension_distance);
      float wrap_distance = dim->Max - dim->Min - dimension_distance;
      dimension_distance = std::min(dimension_distance, wrap_distance);
    }
 
    total_distance += dimension_distance * dimension_distance;
  }
  return total_distance;
}
 
float ComputeDistance(int k, PARAM_DESC *dim, float p1[], float p2[]) {
  return std::sqrt(DistanceSquared(k, dim, p1, p2));
}
 
/*---------------------------------------------------------------------------*/
bool KDTreeSearch::BoxIntersectsSearch(float *lower, float *upper) {
  float *query = query_point_;
  // Compute the sum in higher precision.
  double total_distance = 0.0;
  double radius_squared =
      static_cast<double>(results_.max_insertable_key()) * results_.max_insertable_key();
  PARAM_DESC *dim = &tree_->KeyDesc[0];
 
  for (int i = tree_->KeySize; i > 0; i--, dim++, query++, lower++, upper++) {
    if (dim->NonEssential) {
      continue;
    }
 
    float dimension_distance;
    if (*query < *lower) {
      dimension_distance = *lower - *query;
    } else if (*query > *upper) {
      dimension_distance = *query - *upper;
    } else {
      dimension_distance = 0;
    }
 
    /* if this dimension is circular - check wraparound distance */
    if (dim->Circular) {
      float wrap_distance = FLT_MAX;
      if (*query < *lower) {
        wrap_distance = *query + dim->Max - dim->Min - *upper;
      } else if (*query > *upper) {
        wrap_distance = *lower - (*query - (dim->Max - dim->Min));
      }
      dimension_distance = std::min(dimension_distance, wrap_distance);
    }
 
    total_distance += static_cast<double>(dimension_distance) * dimension_distance;
    if (total_distance >= radius_squared) {
      return false;
    }
  }
  return true;
}
 
/*---------------------------------------------------------------------------*/
void Walk(KDTREE *tree, kdwalk_proc action, ClusteringContext *context, KDNODE *sub_tree, int32_t level) {
  (*action)(context, sub_tree->Data, level);
  if (sub_tree->Left != nullptr) {
    Walk(tree, action, context, sub_tree->Left, NextLevel(tree, level));
  }
  if (sub_tree->Right != nullptr) {
    Walk(tree, action, context, sub_tree->Right, NextLevel(tree, level));
  }
}
 
void InsertNodes(KDTREE *tree, KDNODE *nodes) {
  if (nodes == nullptr) {
    return;
  }
 
  KDStore(tree, nodes->Key, nodes->Data);
  InsertNodes(tree, nodes->Left);
  InsertNodes(tree, nodes->Right);
}
 
} // namespace tesseract