Repository: begeekmyfriend/kdtree Branch: master Commit: bbf2443e4a61 Files: 12 Total size: 43.8 KB Directory structure: gitextract_hi41fn0z/ ├── .gitignore ├── CMakeLists.txt ├── LICENSE ├── Makefile ├── README.md ├── kd_search.m ├── kdtree.c ├── kdtree.h ├── kdtree_bench.c ├── sample.cpp ├── sort_by_r.m └── test.m ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ # Object files *.o *.ko *.obj *.elf # Precompiled Headers *.gch *.pch # Libraries *.lib *.a *.la *.lo # Shared objects (inc. Windows DLLs) *.dll *.so *.so.* *.dylib # Executables *.exe *.out *.app *.i*86 *.x86_64 *.hex # Debug files *.dSYM/ *.su ================================================ FILE: CMakeLists.txt ================================================ cmake_minimum_required(VERSION 2.8) find_package(OpenCV REQUIRED) include_directories(${OpenCV_INCLUDE_DIRS}) if (CMAKE_COMPILER_IS_GNUCXX) set(CMAKE_CXX_FLAGS "-std=c++11 -Wall -O3") endif (CMAKE_COMPILER_IS_GNUCXX) file(GLOB srcs *.c *.cpp *.h*) add_executable(kdtree ${srcs}) target_link_libraries(kdtree ${OpenCV_LIBS}) ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2017 Leo Ma Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: Makefile ================================================ # CMAKE generated file: DO NOT EDIT! # Generated by "Unix Makefiles" Generator, CMake Version 3.26 # Default target executed when no arguments are given to make. default_target: all .PHONY : default_target # Allow only one "make -f Makefile2" at a time, but pass parallelism. .NOTPARALLEL: #============================================================================= # Special targets provided by cmake. # Disable implicit rules so canonical targets will work. .SUFFIXES: # Disable VCS-based implicit rules. % : %,v # Disable VCS-based implicit rules. % : RCS/% # Disable VCS-based implicit rules. % : RCS/%,v # Disable VCS-based implicit rules. % : SCCS/s.% # Disable VCS-based implicit rules. % : s.% .SUFFIXES: .hpux_make_needs_suffix_list # Command-line flag to silence nested $(MAKE). $(VERBOSE)MAKESILENT = -s #Suppress display of executed commands. $(VERBOSE).SILENT: # A target that is always out of date. cmake_force: .PHONY : cmake_force #============================================================================= # Set environment variables for the build. # The shell in which to execute make rules. SHELL = /bin/sh # The CMake executable. CMAKE_COMMAND = /home/zhen/.local/lib/python3.10/site-packages/cmake/data/bin/cmake # The command to remove a file. RM = /home/zhen/.local/lib/python3.10/site-packages/cmake/data/bin/cmake -E rm -f # Escaping for special characters. EQUALS = = # The top-level source directory on which CMake was run. CMAKE_SOURCE_DIR = /home/zhen/work_wls/kdtree # The top-level build directory on which CMake was run. CMAKE_BINARY_DIR = /home/zhen/work_wls/kdtree #============================================================================= # Targets provided globally by CMake. # Special rule for the target edit_cache edit_cache: @$(CMAKE_COMMAND) -E cmake_echo_color --switch=$(COLOR) --cyan "No interactive CMake dialog available..." /home/zhen/.local/lib/python3.10/site-packages/cmake/data/bin/cmake -E echo No\ interactive\ CMake\ dialog\ available. .PHONY : edit_cache # Special rule for the target edit_cache edit_cache/fast: edit_cache .PHONY : edit_cache/fast # Special rule for the target rebuild_cache rebuild_cache: @$(CMAKE_COMMAND) -E cmake_echo_color --switch=$(COLOR) --cyan "Running CMake to regenerate build system..." /home/zhen/.local/lib/python3.10/site-packages/cmake/data/bin/cmake --regenerate-during-build -S$(CMAKE_SOURCE_DIR) -B$(CMAKE_BINARY_DIR) .PHONY : rebuild_cache # Special rule for the target rebuild_cache rebuild_cache/fast: rebuild_cache .PHONY : rebuild_cache/fast # The main all target all: cmake_check_build_system $(CMAKE_COMMAND) -E cmake_progress_start /home/zhen/work_wls/kdtree/CMakeFiles /home/zhen/work_wls/kdtree//CMakeFiles/progress.marks $(MAKE) $(MAKESILENT) -f CMakeFiles/Makefile2 all $(CMAKE_COMMAND) -E cmake_progress_start /home/zhen/work_wls/kdtree/CMakeFiles 0 .PHONY : all # The main clean target clean: $(MAKE) $(MAKESILENT) -f CMakeFiles/Makefile2 clean .PHONY : clean # The main clean target clean/fast: clean .PHONY : clean/fast # Prepare targets for installation. preinstall: all $(MAKE) $(MAKESILENT) -f CMakeFiles/Makefile2 preinstall .PHONY : preinstall # Prepare targets for installation. preinstall/fast: $(MAKE) $(MAKESILENT) -f CMakeFiles/Makefile2 preinstall .PHONY : preinstall/fast # clear depends depend: $(CMAKE_COMMAND) -S$(CMAKE_SOURCE_DIR) -B$(CMAKE_BINARY_DIR) --check-build-system CMakeFiles/Makefile.cmake 1 .PHONY : depend #============================================================================= # Target rules for targets named kdtree # Build rule for target. kdtree: cmake_check_build_system $(MAKE) $(MAKESILENT) -f CMakeFiles/Makefile2 kdtree .PHONY : kdtree # fast build rule for target. kdtree/fast: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/build .PHONY : kdtree/fast kdtree.o: kdtree.c.o .PHONY : kdtree.o # target to build an object file kdtree.c.o: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/kdtree.c.o .PHONY : kdtree.c.o kdtree.i: kdtree.c.i .PHONY : kdtree.i # target to preprocess a source file kdtree.c.i: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/kdtree.c.i .PHONY : kdtree.c.i kdtree.s: kdtree.c.s .PHONY : kdtree.s # target to generate assembly for a file kdtree.c.s: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/kdtree.c.s .PHONY : kdtree.c.s sample.o: sample.cpp.o .PHONY : sample.o # target to build an object file sample.cpp.o: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/sample.cpp.o .PHONY : sample.cpp.o sample.i: sample.cpp.i .PHONY : sample.i # target to preprocess a source file sample.cpp.i: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/sample.cpp.i .PHONY : sample.cpp.i sample.s: sample.cpp.s .PHONY : sample.s # target to generate assembly for a file sample.cpp.s: $(MAKE) $(MAKESILENT) -f CMakeFiles/kdtree.dir/build.make CMakeFiles/kdtree.dir/sample.cpp.s .PHONY : sample.cpp.s # Help Target help: @echo "The following are some of the valid targets for this Makefile:" @echo "... all (the default if no target is provided)" @echo "... clean" @echo "... depend" @echo "... edit_cache" @echo "... rebuild_cache" @echo "... kdtree" @echo "... kdtree.o" @echo "... kdtree.i" @echo "... kdtree.s" @echo "... sample.o" @echo "... sample.i" @echo "... sample.s" .PHONY : help #============================================================================= # Special targets to cleanup operation of make. # Special rule to run CMake to check the build system integrity. # No rule that depends on this can have commands that come from listfiles # because they might be regenerated. cmake_check_build_system: $(CMAKE_COMMAND) -S$(CMAKE_SOURCE_DIR) -B$(CMAKE_BINARY_DIR) --check-build-system CMakeFiles/Makefile.cmake 0 .PHONY : cmake_check_build_system ================================================ FILE: README.md ================================================ # kdtree This is a (nearly absolute) balanced kdtree for fast kNN search. It does not support dynamic insertion and removal. Actually we adopt quick sort to rebuild the whole tree after changes of nodes. We cache the added or the deleted nodes which will not be actually mapped into the tree until the rebuild method to be invoked. The good thing is we can always keep the tree balanced, and the bad thing is we have to wait some time for the finish of tree rebuild. Moreover, duplicated samples are allowed to be added. The thought of the implementation is posted [here](https://www.joinquant.com/post/2843). ## how to debug $ mkdir build $ cd build $ cmake .. $ make $ ./kdtree 1234 # you can try different seed number for random sample generations Image ================================================ FILE: kd_search.m ================================================ function [nearest,nearestIndex] = kd_search(rootIndex,Tree,target, k) %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% % % 2018/8/6 % kdڽ % rootIndexݵ㼯kdĸڵ±(cellеλã % targetĿ % Treekdcell % nearst㼯оĿĵ± % nearst㼯оĿĵ %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% currentNode = rootIndex; %ǰڵ±꣬Ӹڵ㿪ʼ currentNode = search_down(Tree,currentNode,target);% currentNodeײ Tree{currentNode}.visited = 1; currentNearest = Tree{currentNode}.val; % ǰ currentNearestDist = norm(currentNearest-target); % ǰ currentNearestIndex = currentNode; kNearestCandidate = [currentNearestDist, currentNearest', currentNearestIndex]; while Tree{currentNode}.isRoot == 0 isLeft = Tree{currentNode}.isLeft; %ǰڵӱ־ currentNode = Tree{Tree{currentNode}.parent}.index; if Tree{currentNode}.visited == 0 Tree{currentNode}.visited = 1;%Ϊѷ temp = norm(Tree{currentNode}.val-target); if temp=k kNearestCandidate(k+1:end,:) = []; end end temp = abs(Tree{currentNode}.val(Tree{currentNode}.r)-target(Tree{currentNode}.r)); %뵱ǰָߵľ if temp=k kNearestCandidate(k+1:end,:) = []; end end end else if Tree{currentNode}.hasLeft == 1 %ǰڵҺӣҸ׽ڵӣ׽ڵ currentNode = Tree{Tree{currentNode}.left}.index; currentNode = search_down(Tree,currentNode,target); Tree{currentNode}.visited = 1;%Ϊѷ temp = norm(target - Tree{currentNode}.val); if temp=k kNearestCandidate(k+1:end,:) = []; end end end end end end end nearest = kNearestCandidate(:,2:end-1); nearestIndex = kNearestCandidate(:,end); ================================================ FILE: kdtree.c ================================================ /* * Copyright (C) 2017, Leo Ma */ #include #include #include #include #include #include "kdtree.h" static inline int is_leaf(struct kdnode *node) { return node->left == node->right; } static inline void swap(long *a, long *b) { long tmp = *a; *a = *b; *b = tmp; } static inline double square(double d) { return d * d; } static inline double distance(double *c1, double *c2, int dim) { double distance = 0; while (dim-- > 0) { distance += square(*c1++ - *c2++); } return distance; } static inline double knn_max(struct kdtree *tree) { return tree->knn_list_head.prev->distance; } static inline double D(struct kdtree *tree, long index, int r) { return tree->coord_table[index][r]; } static inline int kdnode_passed(struct kdtree *tree, struct kdnode *node) { return node != NULL ? tree->coord_passed[node->coord_index] : 1; } static inline int knn_search_on(struct kdtree *tree, int k, double value, double target) { return tree->knn_num < k || square(target - value) < knn_max(tree); } static inline void coord_index_reset(struct kdtree *tree) { long i; for (i = 0; i < tree->capacity; i++) { tree->coord_indexes[i] = i; } } static inline void coord_table_reset(struct kdtree *tree) { long i; for (i = 0; i < tree->capacity; i++) { tree->coord_table[i] = tree->coords + i * tree->dim; } } static inline void coord_deleted_reset(struct kdtree *tree) { memset(tree->coord_deleted, 0, tree->capacity); } static inline void coord_passed_reset(struct kdtree *tree) { memset(tree->coord_passed, 0, tree->capacity); } static void coord_dump_all(struct kdtree *tree) { long i, j; for (i = 0; i < tree->count; i++) { long index = tree->coord_indexes[i]; double *coord = tree->coord_table[index]; printf("("); for (j = 0; j < tree->dim; j++) { if (j != tree->dim - 1) { printf("%.2f,", coord[j]); } else { printf("%.2f)\n", coord[j]); } } } } static void coord_dump_by_indexes(struct kdtree *tree, long low, long high, int r) { long i; printf("r=%d:", r); for (i = 0; i <= high; i++) { if (i < low) { printf("%8s", " "); } else { long index = tree->coord_indexes[i]; printf("%8.2f", tree->coord_table[index][r]); } } printf("\n"); } static void quicksort(struct kdtree *tree, long lo, long hi, int r) { long i, j, pivot, *indexes; double *p; if (lo >= hi) { return; } i = lo; j = hi; indexes = tree->coord_indexes; pivot = indexes[lo]; while (i < j) { while (i < j && D(tree, indexes[j], r) >= D(tree, pivot, r)) j--; /* Loop invariant: nums[j] > pivot or i == j */ indexes[i] = indexes[j]; while (i < j && D(tree, indexes[i], r) <= D(tree, pivot, r)) i++; /* Loop invariant: nums[i] < pivot or i == j */ indexes[j] = indexes[i]; } /* Loop invariant: i == j */ indexes[i] = pivot; quicksort(tree, lo, i - 1, r); quicksort(tree, i + 1, hi, r); } static struct kdnode *kdnode_alloc(double *coord, long index, int r) { struct kdnode *node = malloc(sizeof(*node)); if (node != NULL) { memset(node, 0, sizeof(*node)); node->coord = coord; node->coord_index = index; node->r = r; } return node; } static void kdnode_free(struct kdnode *node) { free(node); } static int coord_cmp(double *c1, double *c2, int dim) { int i; double ret; for (i = 0; i < dim; i++) { ret = *c1++ - *c2++; if (fabs(ret) >= DBL_EPSILON) { return ret > 0 ? 1 : -1; } } if (fabs(ret) < DBL_EPSILON) { return 0; } else { return ret > 0 ? 1 : -1; } } static void knn_list_add(struct kdtree *tree, struct kdnode *node, double distance) { if (node == NULL) return; struct knn_list *head = &tree->knn_list_head; struct knn_list *p = head->prev; if (tree->knn_num == 1) { if (p->distance > distance) { p = p->prev; } } else { while (p != head && p->distance > distance) { p = p->prev; } } if (p == head || coord_cmp(p->node->coord, node->coord, tree->dim)) { struct knn_list *log = malloc(sizeof(*log)); if (log != NULL) { log->node = node; log->distance = distance; log->prev = p; log->next = p->next; p->next->prev = log; p->next = log; tree->knn_num++; } } } static void knn_list_adjust(struct kdtree *tree, struct kdnode *node, double distance) { if (node == NULL) return; struct knn_list *head = &tree->knn_list_head; struct knn_list *p = head->prev; if (tree->knn_num == 1) { if (p->distance > distance) { p = p->prev; } } else { while (p != head && p->distance > distance) { p = p->prev; } } if (p == head || coord_cmp(p->node->coord, node->coord, tree->dim)) { struct knn_list *log = head->prev; /* Replace the original max one */ log->node = node; log->distance = distance; /* Remove from the max position */ head->prev = log->prev; log->prev->next = head; /* insert as a new one */ log->prev = p; log->next = p->next; p->next->prev = log; p->next = log; } } static void knn_list_clear(struct kdtree *tree) { if (tree->knn_num > 0) { struct knn_list *head = &tree->knn_list_head; struct knn_list *p = head->next; while (p != head) { struct knn_list *prev = p; p = p->next; free(prev); } tree->knn_num = 0; } tree->knn_list_head.next = &tree->knn_list_head; tree->knn_list_head.prev = &tree->knn_list_head; tree->knn_list_head.node = NULL; tree->knn_list_head.distance = 0; } static void resize(struct kdtree *tree) { tree->capacity *= 2; tree->coords = realloc(tree->coords, tree->dim * sizeof(double) * tree->capacity); tree->coord_table = realloc(tree->coord_table, sizeof(double *) * tree->capacity); tree->coord_indexes = realloc(tree->coord_indexes, sizeof(long) * tree->capacity); tree->coord_deleted = realloc(tree->coord_deleted, sizeof(char) * tree->capacity); tree->coord_passed = realloc(tree->coord_passed, sizeof(char) * tree->capacity); coord_table_reset(tree); coord_index_reset(tree); coord_deleted_reset(tree); coord_passed_reset(tree); } static void kdnode_dump(struct kdnode *node, int dim) { int i; if (node->coord != NULL) { printf("("); for (i = 0; i < dim; i++) { if (i != dim - 1) { printf("%.2f,", node->coord[i]); } else { printf("%.2f)\n", node->coord[i]); } } } else { printf("(none)\n"); } } void kdtree_insert(struct kdtree *tree, double *coord) { if (tree->count + 1 > tree->capacity) { resize(tree); } memcpy(tree->coord_table[tree->count++], coord, tree->dim * sizeof(double)); } static void knn_pickup(struct kdtree *tree, struct kdnode *node, double *target, int k) { double dist = distance(node->coord, target, tree->dim); if (tree->knn_num < k) { knn_list_add(tree, node, dist); } else { if (dist < knn_max(tree)) { knn_list_adjust(tree, node, dist); } else if (fabs(dist - knn_max(tree)) < DBL_EPSILON) { knn_list_add(tree, node, dist); } } } static void kdtree_search_recursive(struct kdtree *tree, struct kdnode *node, double *target, int k, int *pickup) { if (node == NULL || kdnode_passed(tree, node)) { return; } int r = node->r; if (is_leaf(node)) { *pickup = 1; } else { if (target[r] <= node->coord[r]) { kdtree_search_recursive(tree, node->left, target, k, pickup); if ((*pickup && knn_search_on(tree, k, node->coord[r], target[r])) || node->left == NULL) { kdtree_search_recursive(tree, node->right, target, k, pickup); } } else { kdtree_search_recursive(tree, node->right, target, k, pickup); if ((*pickup && knn_search_on(tree, k, node->coord[r], target[r])) || node->right == NULL) { kdtree_search_recursive(tree, node->left, target, k, pickup); } } } /* back track and pick up */ if (*pickup) { tree->coord_passed[node->coord_index] = 1; knn_pickup(tree, node, target, k); } } void kdtree_knn_search(struct kdtree *tree, double *target, int k) { if (k > 0) { int pickup = 0; knn_list_clear(tree); coord_passed_reset(tree); kdtree_search_recursive(tree, tree->root, target, k, &pickup); } } void kdtree_delete(struct kdtree *tree, double *coord) { int r = 0; struct kdnode *node = tree->root; struct kdnode *parent = node; while (node != NULL) { if (node->coord == NULL) { if (parent->right->coord == NULL) { break; } else { node = parent->right; continue; } } if (coord[r] < node->coord[r]) { parent = node; node = node->left; } else if (coord[r] > node->coord[r]) { parent = node; node = node->right; } else { int ret = coord_cmp(coord, node->coord, tree->dim); if (ret < 0) { parent = node; node = node->left; } else if (ret > 0) { parent = node; node = node->right; } else { node->coord = NULL; break; } } r = (r + 1) % tree->dim; } } static void kdnode_build(struct kdtree *tree, struct kdnode **nptr, int r, long low, long high) { /* BST preorder constructure */ if (low == high) { long index = tree->coord_indexes[low]; *nptr = kdnode_alloc(tree->coord_table[index], index, r); } else if (low < high) { /* Sort and fetch the median to build a balanced BST */ quicksort(tree, low, high, r); long median = low + (high - low) / 2; long median_index = tree->coord_indexes[median]; struct kdnode *node = *nptr = kdnode_alloc(tree->coord_table[median_index], median_index, r); r = (r + 1) % tree->dim; kdnode_build(tree, &node->left, r, low, median - 1); kdnode_build(tree, &node->right, r, median + 1, high); } } static void kdtree_build(struct kdtree *tree) { kdnode_build(tree, &tree->root, 0, 0, tree->count - 1); } void kdtree_rebuild(struct kdtree *tree) { long i, j; size_t size_of_coord = tree->dim * sizeof(double); for (i = 0, j = 0; j < tree->count; i++, j++) { while (j < tree->count && tree->coord_deleted[j]) { j++; } if (i != j && j < tree->count) { memcpy(tree->coord_table[i], tree->coord_table[j], size_of_coord); tree->coord_deleted[i] = 0; } } tree->count = i; coord_index_reset(tree); kdtree_build(tree); } struct kdtree *kdtree_init(int dim) { struct kdtree *tree = malloc(sizeof(*tree)); if (tree != NULL) { tree->root = NULL; tree->dim = dim; tree->count = 0; tree->capacity = 65536; tree->knn_num = 0; tree->coords = malloc(dim * sizeof(double) * tree->capacity); tree->coord_table = malloc(sizeof(double *) * tree->capacity); tree->coord_indexes = malloc(sizeof(long) * tree->capacity); tree->coord_deleted = malloc(sizeof(char) * tree->capacity); tree->coord_passed = malloc(sizeof(char) * tree->capacity); coord_index_reset(tree); coord_table_reset(tree); coord_deleted_reset(tree); coord_passed_reset(tree); } return tree; } static void kdnode_destroy(struct kdnode *node) { if (node == NULL) return; kdnode_destroy(node->left); kdnode_destroy(node->right); kdnode_free(node); } void kdtree_destroy(struct kdtree *tree) { kdnode_destroy(tree->root); knn_list_clear(tree); free(tree->coords); free(tree->coord_table); free(tree->coord_indexes); free(tree->coord_deleted); free(tree->coord_passed); free(tree); } #define _KDTREE_DEBUG #ifdef _KDTREE_DEBUG struct kdnode_backlog { struct kdnode *node; int next_sub_idx; }; void kdtree_dump(struct kdtree *tree) { int level = 0; struct kdnode *node = tree->root; struct kdnode_backlog nbl, *p_nbl = NULL; struct kdnode_backlog nbl_stack[KDTREE_MAX_LEVEL]; struct kdnode_backlog *top = nbl_stack; for (; ;) { if (node != NULL) { /* Fetch the pop-up backlogged node's sub-id. * If not backlogged, fetch the first sub-id. */ int sub_idx = p_nbl != NULL ? p_nbl->next_sub_idx : KDTREE_RIGHT_INDEX; /* Backlog should be left in next loop */ p_nbl = NULL; /* Backlog the node */ if (is_leaf(node) || sub_idx == KDTREE_LEFT_INDEX) { top->node = NULL; top->next_sub_idx = KDTREE_RIGHT_INDEX; } else { top->node = node; top->next_sub_idx = KDTREE_LEFT_INDEX; } top++; level++; /* Draw lines as long as sub_idx is the first one */ if (sub_idx == KDTREE_RIGHT_INDEX) { int i; for (i = 1; i < level; i++) { if (i == level - 1) { printf("%-8s", "+-------"); } else { if (nbl_stack[i - 1].node != NULL) { printf("%-8s", "|"); } else { printf("%-8s", " "); } } } kdnode_dump(node, tree->dim); } /* Move down according to sub_idx */ node = sub_idx == KDTREE_LEFT_INDEX ? node->left : node->right; } else { p_nbl = top == nbl_stack ? NULL : --top; if (p_nbl == NULL) { /* End of traversal */ break; } node = p_nbl->node; level--; } } } #endif // (Includes and struct definitions from kdtree.h) // (kdnode_dump helper function from original kdtree.c) /* * Helper struct for our new Root-Left-Right stack. * We need to know the node and its level. */ struct kdnode_dump_item { struct kdnode *node; int level; }; /* * This is the modified kdtree_dump function. * It uses a standard iterative Pre-order (Root-Left-Right) traversal. * * It manages the line-art state explicitly using a * 'level_states' array. */ void kdtree_dump_simple(struct kdtree *tree) { if (tree->root == NULL) { printf("Tree is empty.\n"); return; } /* * 1. Create our stack */ struct kdnode_dump_item stack[KDTREE_MAX_LEVEL]; int top = -1; /* * 2. Create state for line-drawing * This array tracks whether to draw a '|' at a given level. * 1 = draw '|', 0 = draw ' ' */ int level_states[KDTREE_MAX_LEVEL]; memset(level_states, 0, sizeof(level_states)); /* * 3. Push the root node to start */ top++; stack[top].node = tree->root; stack[top].level = 0; while (top != -1) { /* * Pop a node from the stack */ struct kdnode_dump_item current_item = stack[top--]; struct kdnode *node = current_item.node; int level = current_item.level; /* --- This is the "Visit" step --- */ /* Print the prefix lines ('|' or ' ') */ int i; for (i = 0; i < level; i++) { if (level_states[i] == 2) { printf("%-8s", "|"); } else { printf("%-8s", " "); } } /* Print the node itself */ printf("%-8s", "+-------"); kdnode_dump(node, tree->dim); /* --- End Visit --- */ level_states[level]++; /* * 4. Update line-drawing state for the *next* level * If we have a left child, the right child (if it exists) * will need a '|' at this level. */ if (node->left != NULL && node->right != NULL) { level_states[level+1] = 1; } else { level_states[level+1] = 0; } /* * 5. Push children onto the stack (Root-Left-Right) * We push RIGHT first, then LEFT. */ if (node->left != NULL) { top++; stack[top].node = node->left; stack[top].level = level + 1; } if (node->right != NULL) { top++; stack[top].node = node->right; stack[top].level = level + 1; } } } ================================================ FILE: kdtree.h ================================================ /* * Copyright (C) 2017, Leo Ma */ #ifndef _KD_TREE_H #define _KD_TREE_H /* ADD THIS LINE */ #ifdef __cplusplus extern "C" { #endif #define KDTREE_MAX_LEVEL 64 #define KDTREE_LEFT_INDEX 0 #define KDTREE_RIGHT_INDEX 1 typedef struct knn_list { struct kdnode *node; double distance; struct knn_list *prev; struct knn_list *next; } knn_list_t; typedef struct kdnode { long coord_index; double *coord; struct kdnode *left; struct kdnode *right; int r; } kdnode_t; typedef struct kdtree { struct kdnode *root; size_t count; size_t capacity; double *coords; double **coord_table; long *coord_indexes; unsigned char *coord_deleted; unsigned char *coord_passed; struct knn_list knn_list_head; int dim; int knn_num; } kdtree_t; struct kdtree *kdtree_init(int dim); void kdtree_insert(struct kdtree *tree, double *coord); void kdtree_rebuild(struct kdtree *tree); void kdtree_knn_search(struct kdtree *tree, double *coord, int k); void kdtree_destroy(struct kdtree *tree); void kdtree_dump(struct kdtree *tree); void kdtree_dump_simple(struct kdtree *tree); /* ADD THESE LINES */ #ifdef __cplusplus } #endif #endif /* _KD_TREE_H */ ================================================ FILE: kdtree_bench.c ================================================ #include #include #include #include #include "kdtree.h" #define N 1024 * 1024 static inline double rd(void) { return (double) rand() / RAND_MAX * 20 - 10; } static void kdtree_knn_dump(struct kdtree *tree) { int i; struct knn_list *p = tree->knn_list_head.next; while (p != &tree->knn_list_head) { putchar('('); for (i = 0; i < tree->dim; i++) { if (i == tree->dim - 1) { printf("%.2lf) Distance:%lf\n", p->node->coord[i], sqrt(p->distance)); } else { printf("%.2lf, ", p->node->coord[i]); } } p = p->next; } } int main(void) { int i, j, dim = 2; struct timespec start, end; struct kdtree *tree = kdtree_init(dim); if (tree == NULL) { exit(-1); } double sample1[] = { 6.27, 5.50 }; double sample2[] = { 1.24, -2.86 }; double sample3[] = { 17.05, -12.79 }; double sample4[] = { -6.88, -5.40 }; double sample5[] = { -2.96, -0.50 }; double sample6[] = { 7.75, -22.68 }; double sample7[] = { 10.80, -5.03 }; double sample8[] = { -4.60, -10.55 }; double sample9[] = { -4.96, 12.61 }; double sample10[] = { 1.75, 12.26 }; double sample11[] = { 15.31, -13.16 }; double sample12[] = { 7.83, 15.70 }; double sample13[] = { 14.63, -0.35 }; kdtree_insert(tree, sample1); kdtree_insert(tree, sample2); kdtree_insert(tree, sample3); kdtree_insert(tree, sample4); kdtree_insert(tree, sample5); kdtree_insert(tree, sample6); kdtree_insert(tree, sample7); kdtree_insert(tree, sample8); kdtree_insert(tree, sample9); kdtree_insert(tree, sample10); kdtree_insert(tree, sample11); kdtree_insert(tree, sample12); kdtree_insert(tree, sample13); kdtree_rebuild(tree); kdtree_dump(tree); int k = 3; double target[] = { -1, -5 }; kdtree_knn_search(tree, target, k); printf("%d nearest neighbors of sample(", k); for (i = 0; i < dim; i++) { if (i == dim - 1) { printf("%.2lf):\n", target[i]); } else { printf("%.2lf, ", target[i]); } } kdtree_knn_dump(tree); double target1[] = { -8, -7 }; kdtree_knn_search(tree, target1, k); printf("%d nearest neighbors of sample(", k); for (i = 0; i < dim; i++) { if (i == dim - 1) { printf("%.2lf):\n", target1[i]); } else { printf("%.2lf, ", target1[i]); } } kdtree_knn_dump(tree); kdtree_destroy(tree); /* Performance test */ printf("\n>>> Performance test: kNN search for %d samples\n\n", N); dim = 12; tree = kdtree_init(dim); if (tree == NULL) { exit(-1); } /* Insert test */ printf("Add %d nodes...\n", N); srandom(time(NULL)); clock_gettime(CLOCK_MONOTONIC, &start); for (i = 0; i < N; i++) { double *sample = malloc(dim * sizeof(double)); for (j = 0; j < dim; j++) { sample[j] = rd(); } kdtree_insert(tree, sample); } clock_gettime(CLOCK_MONOTONIC, &end); printf("time span: %ldms\n", (end.tv_sec - start.tv_sec)*1000 + (end.tv_nsec - start.tv_nsec)/1000000); /* Build test */ printf("Build KD tree...\n"); clock_gettime(CLOCK_MONOTONIC, &start); kdtree_rebuild(tree); clock_gettime(CLOCK_MONOTONIC, &end); printf("time span: %ldms\n", (end.tv_sec - start.tv_sec)*1000 + (end.tv_nsec - start.tv_nsec)/1000000); /* Search test */ k = 20; srandom(time(NULL)); double *t = malloc(dim * sizeof(double)); for (i = 0; i < dim; i++) { t[i] = rd(); } printf("Search KD tree...\n"); srandom(time(NULL)); clock_gettime(CLOCK_MONOTONIC, &start); kdtree_knn_search(tree, t, k); clock_gettime(CLOCK_MONOTONIC, &end); printf("time span: %ldms\n", (end.tv_sec - start.tv_sec)*1000 + (end.tv_nsec - start.tv_nsec)/1000000); printf("%d nearest neighbors of sample(", k); for (i = 0; i < dim; i++) { if (i == dim - 1) { printf("%.2lf):\n", t[i]); } else { printf("%.2lf, ", t[i]); } } kdtree_knn_dump(tree); /* Destroy test */ printf("Destroy KD tree...\n"); srandom(time(NULL)); clock_gettime(CLOCK_MONOTONIC, &start); kdtree_destroy(tree); clock_gettime(CLOCK_MONOTONIC, &end); printf("time span: %ldms\n", (end.tv_sec - start.tv_sec)*1000 + (end.tv_nsec - start.tv_nsec)/1000000); return 0; } ================================================ FILE: sample.cpp ================================================ #include #include #include #include #include "kdtree.h" // user-defined point type // inherits std::array in order to use operator[] class MyPoint : public std::array { public: // dimension of space (or "k" of k-d tree) // KDTree class accesses this member static const int DIM = 2; // the constructors MyPoint() {} MyPoint(double x, double y) { (*this)[0] = x; (*this)[1] = y; } // conversion to OpenCV Point2d operator cv::Point2d() const { return cv::Point2d((*this)[0], (*this)[1]); } }; void kdtree_knn_dump(struct kdtree *tree, double *candidates, int k) { int i, j = 0; struct knn_list *p = tree->knn_list_head.next; const int dim = tree->dim; printf("The nearest %d samples are as follows:\n", k); while (p != &tree->knn_list_head) { putchar('('); for (i = 0; i < tree->dim; i++) { if (i == tree->dim - 1) { printf("%.2lf) Distance:%lf\n", p->node->coord[i], sqrt(p->distance)); } else { printf("%.2lf, ", p->node->coord[i]); } candidates[j*dim+i] = p->node->coord[i]; } p = p->next; j++; } } int main(int argc, char **argv) { const int seed = argc > 1 ? std::stoi(argv[1]) : 0; srand(seed); // generate space const int width = 500; const int height = 500; cv::Mat img = cv::Mat::zeros(cv::Size(width, height), CV_8UC3); // generate points const int npoints = 1000; std::vector points(npoints); for (int i = 0; i < npoints; i++) { const int x = rand() % width; const int y = rand() % height; points[i] = MyPoint(x, y); } for (const auto& pt : points) cv::circle(img, cv::Point2d(pt), 1, cv::Scalar(0, 255, 255), -1); // build k-d tree int dim = 2; struct kdtree *tree = kdtree_init(dim); for (int i = 0; i < npoints; i++) { // 1. Declare an array to hold the coordinates double my_coord[] = {points[i][0], points[i][1]}; // 2. Pass the array (which acts as a pointer) to the function kdtree_insert(tree, my_coord); } kdtree_rebuild(tree); // generate query (center of the space) const int k = 100; const MyPoint query(0.5 * width, 0.5 * height); cv::circle(img, cv::Point2d(query), 1, cv::Scalar(0, 0, 255), -1); // nearest neigbor search const cv::Mat I0 = img.clone(); double target1[] = {query[0], query[1]}; kdtree_knn_search(tree, target1, 1); // kdtree_knn_dump(tree, nullptr, k); // const int idx = kdtree.nnSearch(query); // cv::circle(I0, cv::Point2d(points[idx]), 1, cv::Scalar(255, 255, 0), -1); // cv::line(I0, cv::Point2d(query), cv::Point2d(points[idx]), cv::Scalar(0, 0, 255)); // k-nearest neigbors search const cv::Mat I1 = img.clone(); double candidates[k*dim]; kdtree_knn_search(tree, target1, k); kdtree_knn_dump(tree, candidates, k); for (int i = 0; i < k; i++){ cv::circle(I1, cv::Point2d(candidates[i*dim], candidates[i*dim + 1]), 1, cv::Scalar(255, 255, 0), -1); cv::line(I1, cv::Point2d(query), cv::Point2d(candidates[i*dim], candidates[i*dim + 1]), cv::Scalar(0, 0, 255)); } // const std::vector knnIndices = kdtree.knnSearch(query, k); // for (int i : knnIndices) // { // cv::circle(I1, cv::Point2d(points[i]), 1, cv::Scalar(255, 255, 0), -1); // cv::line(I1, cv::Point2d(query), cv::Point2d(points[i]), cv::Scalar(0, 0, 255)); // } // radius search const cv::Mat I2 = img.clone(); // const double radius = 50; // const std::vector radIndices = kdtree.radiusSearch(query, radius); // for (int i : radIndices) // cv::circle(I2, cv::Point2d(points[i]), 1, cv::Scalar(255, 255, 0), -1); // cv::circle(I2, cv::Point2d(query), cvRound(radius), cv::Scalar(0, 0, 255)); // show results cv::imshow("Nearest neigbor search", I0); // cv::imshow("K-nearest neigbors search (k = 10)", I1); cv::imshow("K-nearest neighbors search (k = " + std::to_string(k) + ")", I1); cv::imshow("Radius search (radius = 50)", I2); cv::waitKey(0); return 0; } ================================================ FILE: sort_by_r.m ================================================ function X_sorted = sort_by_r(X,r) %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% % ݵ㼯rά % Xݵ㼯[dim,num] % rά % X_sorted㼯 % ð %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% % num = size(X,2); % for i = 1:num-1 % for j = 1:num-i % if X(r,j+1)