sc2ai-kdtree 0.4.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +7 -0
- data/.github/workflows/ci.yml +16 -0
- data/.gitignore +70 -0
- data/Gemfile +2 -0
- data/LICENSE +20 -0
- data/README.md +5 -0
- data/Rakefile +41 -0
- data/ext/kdtree/extconf.rb +3 -0
- data/ext/kdtree/kdtree.c +505 -0
- data/kdtree.gemspec +22 -0
- data/lib/kdtree.rb +1 -0
- data/test/test_kdtree.rb +148 -0
- metadata +86 -0
checksums.yaml
ADDED
@@ -0,0 +1,7 @@
|
|
1
|
+
---
|
2
|
+
SHA256:
|
3
|
+
metadata.gz: 1b02e78ec134ea010fae56b8e79c189453f5d2e62aa932c553366463f336ef88
|
4
|
+
data.tar.gz: bda73cdbdf514e92593e06738573a1389f4a1f190a29281a2a342dab71ff007e
|
5
|
+
SHA512:
|
6
|
+
metadata.gz: 30b32791ff9d76b5bf24d9b28730dfa72fa41c68b2bd7972511a75de782a48244caccb27e29a371c6373805bfcd7359d4867027763c922a7575625d1f186f562
|
7
|
+
data.tar.gz: 9fb55e05aef40736c61ca61e84d8d21de7df357e82419281feca054920da557af3122794d3d5fd50787f7d5d596497bb1be1a8df24e810764150677940478247
|
@@ -0,0 +1,16 @@
|
|
1
|
+
name: ci
|
2
|
+
on: [push]
|
3
|
+
jobs:
|
4
|
+
test:
|
5
|
+
runs-on: ubuntu-latest
|
6
|
+
strategy:
|
7
|
+
fail-fast: false
|
8
|
+
matrix:
|
9
|
+
ruby: [2.0, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, head]
|
10
|
+
steps:
|
11
|
+
- uses: actions/checkout@v2
|
12
|
+
- uses: ruby/setup-ruby@v1
|
13
|
+
with:
|
14
|
+
ruby-version: ${{ matrix.ruby }}
|
15
|
+
bundler-cache: true # bundle install
|
16
|
+
- run: bundle exec rake
|
data/.gitignore
ADDED
@@ -0,0 +1,70 @@
|
|
1
|
+
*.gem
|
2
|
+
.ruby-version
|
3
|
+
Gemfile.lock
|
4
|
+
lib/*.so
|
5
|
+
lib/kdtree.bundle
|
6
|
+
tmp
|
7
|
+
|
8
|
+
|
9
|
+
|
10
|
+
*.rbc
|
11
|
+
/.config
|
12
|
+
/coverage/
|
13
|
+
/InstalledFiles
|
14
|
+
/pkg/
|
15
|
+
/spec/reports/
|
16
|
+
/spec/examples.txt
|
17
|
+
/test/tmp/
|
18
|
+
/test/version_tmp/
|
19
|
+
/tmp/
|
20
|
+
|
21
|
+
# rspec failure tracking
|
22
|
+
.rspec_status
|
23
|
+
|
24
|
+
# Used by dotenv library to load environment variables.
|
25
|
+
.env
|
26
|
+
|
27
|
+
# Ignore Byebug command history file.
|
28
|
+
.byebug_history
|
29
|
+
|
30
|
+
## Specific to RubyMotion:
|
31
|
+
.dat*
|
32
|
+
.repl_history
|
33
|
+
build/
|
34
|
+
*.bridgesupport
|
35
|
+
build-iPhoneOS/
|
36
|
+
build-iPhoneSimulator/
|
37
|
+
|
38
|
+
## Specific to RubyMotion (use of CocoaPods):
|
39
|
+
#
|
40
|
+
# We recommend against adding the Pods directory to your .gitignore. However
|
41
|
+
# you should judge for yourself, the pros and cons are mentioned at:
|
42
|
+
# https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control
|
43
|
+
#
|
44
|
+
# vendor/Pods/
|
45
|
+
|
46
|
+
## Specific to RubyMine
|
47
|
+
/.idea
|
48
|
+
|
49
|
+
## Documentation cache and generated files:
|
50
|
+
/.yardoc/
|
51
|
+
/_yardoc/
|
52
|
+
/doc/
|
53
|
+
/rdoc/
|
54
|
+
|
55
|
+
## Environment normalization:
|
56
|
+
/.bundle/
|
57
|
+
/vendor/bundle
|
58
|
+
/lib/bundler/man/
|
59
|
+
|
60
|
+
# for a library or gem, you might want to ignore these files since the code is
|
61
|
+
# intended to run in multiple environments; otherwise, check them in:
|
62
|
+
.ruby-gemset
|
63
|
+
|
64
|
+
# unless supporting rvm < 1.11.0 or doing something fancy, ignore this:
|
65
|
+
.rvmrc
|
66
|
+
|
67
|
+
# Used by RuboCop. Remote config files pulled in from inherit_from directive.
|
68
|
+
# .rubocop-https?--*
|
69
|
+
|
70
|
+
|
data/Gemfile
ADDED
data/LICENSE
ADDED
@@ -0,0 +1,20 @@
|
|
1
|
+
Copyright (c) 2012 Adam Doppelt
|
2
|
+
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining
|
4
|
+
a copy of this software and associated documentation files (the
|
5
|
+
"Software"), to deal in the Software without restriction, including
|
6
|
+
without limitation the rights to use, copy, modify, merge, publish,
|
7
|
+
distribute, sublicense, and/or sell copies of the Software, and to
|
8
|
+
permit persons to whom the Software is furnished to do so, subject to
|
9
|
+
the following conditions:
|
10
|
+
|
11
|
+
The above copyright notice and this permission notice shall be
|
12
|
+
included in all copies or substantial portions of the Software.
|
13
|
+
|
14
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
15
|
+
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
16
|
+
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
17
|
+
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
|
18
|
+
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
19
|
+
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
20
|
+
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
data/README.md
ADDED
data/Rakefile
ADDED
@@ -0,0 +1,41 @@
|
|
1
|
+
require "bundler/setup"
|
2
|
+
require "rake/extensiontask"
|
3
|
+
require "rake/testtask"
|
4
|
+
|
5
|
+
# load the spec, we use it below
|
6
|
+
spec = Gem::Specification.load("kdtree.gemspec")
|
7
|
+
|
8
|
+
#
|
9
|
+
# gem
|
10
|
+
#
|
11
|
+
|
12
|
+
task :build do
|
13
|
+
system "gem build --quiet kdtree.gemspec"
|
14
|
+
end
|
15
|
+
|
16
|
+
task install: :build do
|
17
|
+
system "sudo gem install --quiet sc2ai-kdtree-#{spec.version}.gem"
|
18
|
+
end
|
19
|
+
|
20
|
+
task release: :build do
|
21
|
+
system "git tag -a #{spec.version} -m 'Tagging #{spec.version}'"
|
22
|
+
system "git push --tags"
|
23
|
+
system "gem push kdtree-#{spec.version}.gem"
|
24
|
+
end
|
25
|
+
|
26
|
+
#
|
27
|
+
# rake-compiler
|
28
|
+
#
|
29
|
+
|
30
|
+
Rake::ExtensionTask.new("kdtree", spec)
|
31
|
+
|
32
|
+
|
33
|
+
#
|
34
|
+
# testing
|
35
|
+
#
|
36
|
+
|
37
|
+
Rake::TestTask.new(:test) do |test|
|
38
|
+
test.libs << "test"
|
39
|
+
end
|
40
|
+
task test: :compile
|
41
|
+
task default: :test
|
data/ext/kdtree/kdtree.c
ADDED
@@ -0,0 +1,505 @@
|
|
1
|
+
#include "ruby.h"
|
2
|
+
|
3
|
+
//
|
4
|
+
// interface
|
5
|
+
//
|
6
|
+
|
7
|
+
// the tree itself
|
8
|
+
typedef struct kdtree_data
|
9
|
+
{
|
10
|
+
int root;
|
11
|
+
int len;
|
12
|
+
struct kdtree_node *nodes;
|
13
|
+
} kdtree_data;
|
14
|
+
|
15
|
+
// a node in the tree
|
16
|
+
typedef struct kdtree_node
|
17
|
+
{
|
18
|
+
float x, y;
|
19
|
+
int id;
|
20
|
+
int left;
|
21
|
+
int right;
|
22
|
+
} kdtree_node;
|
23
|
+
|
24
|
+
// a result node from kdtree_nearestk0
|
25
|
+
typedef struct kresult {
|
26
|
+
int index;
|
27
|
+
float distance;
|
28
|
+
} kresult;
|
29
|
+
|
30
|
+
// helper macro for digging out our struct
|
31
|
+
#define KDTREEP \
|
32
|
+
struct kdtree_data *kdtreep; \
|
33
|
+
Data_Get_Struct(kdtree, struct kdtree_data, kdtreep);
|
34
|
+
|
35
|
+
// kdtree public methods
|
36
|
+
static VALUE kdtree_alloc(VALUE klass);
|
37
|
+
static void kdtree_free(struct kdtree_data *kdtreep);
|
38
|
+
static VALUE kdtree_initialize(VALUE kdtree, VALUE points);
|
39
|
+
static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y);
|
40
|
+
static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k);
|
41
|
+
static VALUE kdtree_persist(VALUE kdtree, VALUE io);
|
42
|
+
static VALUE kdtree_to_s(VALUE kdtree);
|
43
|
+
|
44
|
+
// kdtree helpers
|
45
|
+
static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth);
|
46
|
+
static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth, int *n_index, float *n_dist);
|
47
|
+
static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth, kresult *k_list, int *k_len, float *k_dist);
|
48
|
+
|
49
|
+
// io helpers
|
50
|
+
static void read_all(VALUE io, void *buf, int len);
|
51
|
+
static void write_all(VALUE io, const void *buf, int len);
|
52
|
+
|
53
|
+
#define KDTREE_MAGIC "KdTr"
|
54
|
+
|
55
|
+
// ids
|
56
|
+
static ID id_read, id_write, id_binmode;
|
57
|
+
|
58
|
+
//
|
59
|
+
// implementation
|
60
|
+
//
|
61
|
+
|
62
|
+
static VALUE kdtree_alloc(VALUE klass)
|
63
|
+
{
|
64
|
+
struct kdtree_data *kdtreep;
|
65
|
+
VALUE obj = Data_Make_Struct(klass, struct kdtree_data, 0, kdtree_free, kdtreep);
|
66
|
+
kdtreep->root = -1;
|
67
|
+
return obj;
|
68
|
+
}
|
69
|
+
|
70
|
+
static void kdtree_free(struct kdtree_data *kdtreep)
|
71
|
+
{
|
72
|
+
if (kdtreep) {
|
73
|
+
free(kdtreep->nodes);
|
74
|
+
}
|
75
|
+
}
|
76
|
+
|
77
|
+
/*
|
78
|
+
* call-seq:
|
79
|
+
* Kdtree.new(points) => kdtree
|
80
|
+
* Kdtree.new(io) => kdtree
|
81
|
+
*
|
82
|
+
* Returns a new <code>Kdtree</code>. To construct a tree, pass an array of
|
83
|
+
* <i>points</i>. Each point should be an array of the form <code>[x, y,
|
84
|
+
* id]</code>, where <i>x</i> and <i>y</i> are floats and <i>id</i> is an
|
85
|
+
* integer. The <i>id</i> is arbitrary and will be returned to you whenever you
|
86
|
+
* search with nearest or nearestk.
|
87
|
+
*
|
88
|
+
* # create a new tree
|
89
|
+
* points = []
|
90
|
+
* points << [47.6, -122.3, 1] # Seattle
|
91
|
+
* points << [40.7, -74.0, 2] # New York
|
92
|
+
* kd = Kdtree.new(points)
|
93
|
+
*
|
94
|
+
* Alternately, you can pass in an <i>IO</i> object containing a persisted
|
95
|
+
* kdtree. This makes it possible to build the tree in advance, persist it, and
|
96
|
+
* start it up quickly later. See persist for more information.
|
97
|
+
*/
|
98
|
+
static VALUE kdtree_initialize(VALUE kdtree, VALUE arg)
|
99
|
+
{
|
100
|
+
KDTREEP;
|
101
|
+
|
102
|
+
if (TYPE(arg) == T_ARRAY) {
|
103
|
+
// init from array of pints
|
104
|
+
VALUE points = arg;
|
105
|
+
int i;
|
106
|
+
kdtreep->len = (int)RARRAY_LEN(points);
|
107
|
+
kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
|
108
|
+
|
109
|
+
for (i = 0; i < RARRAY_LEN(points); ++i) {
|
110
|
+
struct kdtree_node *n = kdtreep->nodes + i;
|
111
|
+
|
112
|
+
VALUE ptr = rb_ary_entry(points, i);
|
113
|
+
VALUE v = rb_check_array_type(ptr);
|
114
|
+
if (NIL_P(v) || RARRAY_LEN(v) != 3) {
|
115
|
+
continue;
|
116
|
+
}
|
117
|
+
n->x = NUM2DBL(rb_ary_entry(v, 0));
|
118
|
+
n->y = NUM2DBL(rb_ary_entry(v, 1));
|
119
|
+
n->id = NUM2INT(rb_ary_entry(v, 2));
|
120
|
+
}
|
121
|
+
|
122
|
+
// now build the tree
|
123
|
+
kdtreep->root = kdtree_build(kdtreep, 0, kdtreep->len, 0);
|
124
|
+
} else if (rb_respond_to(arg, rb_intern("read"))) {
|
125
|
+
VALUE io = arg;
|
126
|
+
char buf[4];
|
127
|
+
if (rb_respond_to(io, id_binmode)) {
|
128
|
+
rb_funcall(io, id_binmode, 0);
|
129
|
+
}
|
130
|
+
|
131
|
+
// check magic
|
132
|
+
read_all(io, buf, 4);
|
133
|
+
if (memcmp(KDTREE_MAGIC, buf, 4) != 0) {
|
134
|
+
rb_raise(rb_eRuntimeError, "wrong magic number in kdtree file");
|
135
|
+
}
|
136
|
+
|
137
|
+
// read start of the struct
|
138
|
+
read_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
|
139
|
+
|
140
|
+
// read the nodes
|
141
|
+
kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
|
142
|
+
read_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
|
143
|
+
} else {
|
144
|
+
rb_raise(rb_eTypeError, "array or IO required to init Kdtree");
|
145
|
+
}
|
146
|
+
|
147
|
+
return kdtree;
|
148
|
+
}
|
149
|
+
|
150
|
+
static int comparex(const void *pa, const void *pb)
|
151
|
+
{
|
152
|
+
float a = ((const struct kdtree_node*)pa)->x;
|
153
|
+
float b = ((const struct kdtree_node*)pb)->x;
|
154
|
+
return (a < b) ? -1 : ((a > b) ? 1 : 0);
|
155
|
+
}
|
156
|
+
|
157
|
+
static int comparey(const void *pa, const void *pb)
|
158
|
+
{
|
159
|
+
float a = ((const struct kdtree_node*)pa)->y;
|
160
|
+
float b = ((const struct kdtree_node*)pb)->y;
|
161
|
+
return (a < b) ? -1 : ((a > b) ? 1 : 0);
|
162
|
+
}
|
163
|
+
|
164
|
+
static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth)
|
165
|
+
{
|
166
|
+
int(*compar)(const void *, const void *);
|
167
|
+
struct kdtree_node *m;
|
168
|
+
int median;
|
169
|
+
if (max <= min) {
|
170
|
+
return -1;
|
171
|
+
}
|
172
|
+
|
173
|
+
// sort nodes from min to max
|
174
|
+
compar = (depth % 2) ? comparex : comparey;
|
175
|
+
qsort(kdtreep->nodes + min, max - min, sizeof(struct kdtree_node), compar);
|
176
|
+
|
177
|
+
median = (min + max) / 2;
|
178
|
+
m = kdtreep->nodes + median;
|
179
|
+
m->left = kdtree_build(kdtreep, min, median, depth + 1);
|
180
|
+
m->right = kdtree_build(kdtreep, median + 1, max, depth + 1);
|
181
|
+
return median;
|
182
|
+
}
|
183
|
+
|
184
|
+
/*
|
185
|
+
* call-seq:
|
186
|
+
* kd.nearest(x, y) => id
|
187
|
+
*
|
188
|
+
* Finds the point closest to <i>x</i>, <i>y</i> and returns the id for that
|
189
|
+
* point. Returns -1 if the tree is empty.
|
190
|
+
*
|
191
|
+
* points = []
|
192
|
+
* points << [47.6, -122.3, 1] # Seattle
|
193
|
+
* points << [40.7, -74.0, 2] # New York
|
194
|
+
* kd = Kdtree.new(points)
|
195
|
+
*
|
196
|
+
* # which city is closest to Portland?
|
197
|
+
* kd.nearest(45.5, -122.8) #=> 1
|
198
|
+
* # which city is closest to Boston?
|
199
|
+
* kd.nearest(42.4, -71.1) #=> 2
|
200
|
+
*/
|
201
|
+
static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y)
|
202
|
+
{
|
203
|
+
int n_index;
|
204
|
+
float n_dist;
|
205
|
+
KDTREEP;
|
206
|
+
|
207
|
+
n_index = -1;
|
208
|
+
n_dist = INT_MAX;
|
209
|
+
|
210
|
+
kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0, &n_index, &n_dist);
|
211
|
+
if (n_index == -1) {
|
212
|
+
return -1;
|
213
|
+
}
|
214
|
+
return INT2NUM((kdtreep->nodes + n_index)->id);
|
215
|
+
}
|
216
|
+
|
217
|
+
static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth, int *n_index, float *n_dist)
|
218
|
+
{
|
219
|
+
struct kdtree_node *n;
|
220
|
+
float ad;
|
221
|
+
int nearer, further;
|
222
|
+
float dx;
|
223
|
+
|
224
|
+
if (i == -1) {
|
225
|
+
return;
|
226
|
+
}
|
227
|
+
|
228
|
+
n = kdtreep->nodes + i;
|
229
|
+
|
230
|
+
ad = (depth % 2) ? (x - n->x) : (y - n->y);
|
231
|
+
|
232
|
+
//
|
233
|
+
// recurse near, and perhaps far as well
|
234
|
+
//
|
235
|
+
|
236
|
+
if (ad <= 0) {
|
237
|
+
nearer = n->left; further = n->right;
|
238
|
+
} else {
|
239
|
+
nearer = n->right; further = n->left;
|
240
|
+
}
|
241
|
+
kdtree_nearest0(kdtreep, nearer, x, y, depth + 1, n_index, n_dist);
|
242
|
+
if (ad * ad < *n_dist) {
|
243
|
+
kdtree_nearest0(kdtreep, further, x, y, depth + 1, n_index, n_dist);
|
244
|
+
}
|
245
|
+
|
246
|
+
//
|
247
|
+
// do we beat the old distance?
|
248
|
+
//
|
249
|
+
|
250
|
+
dx = (x - n->x) * (x - n->x);
|
251
|
+
if (dx < *n_dist) {
|
252
|
+
float d = dx + ((y - n->y) * (y - n->y));
|
253
|
+
if (d < *n_dist) {
|
254
|
+
*n_index = i;
|
255
|
+
*n_dist = d;
|
256
|
+
}
|
257
|
+
}
|
258
|
+
}
|
259
|
+
|
260
|
+
//
|
261
|
+
// nearestK
|
262
|
+
//
|
263
|
+
|
264
|
+
#define MAX_K 255
|
265
|
+
|
266
|
+
/*
|
267
|
+
* call-seq:
|
268
|
+
* kd.nearestk(x, y, k) => array
|
269
|
+
*
|
270
|
+
* Finds the <i>k</i> points closest to <i>x</i>, <i>y</i>. Returns an array of
|
271
|
+
* ids, sorted by distance. Returns an empty array if the tree is empty. Note
|
272
|
+
* that <i>k</i> is capped at 255.
|
273
|
+
*
|
274
|
+
* points = []
|
275
|
+
* points << [47.6, -122.3, 1] # Seattle
|
276
|
+
* points << [45.5, -122.8, 2] # Portland
|
277
|
+
* points << [40.7, -74.0, 3] # New York
|
278
|
+
* kd = Kdtree.new(points)
|
279
|
+
*
|
280
|
+
* # which two cities are closest to San Francisco?
|
281
|
+
* kd.nearestk(34.1, -118.2, 2) #=> [2, 1]
|
282
|
+
*/
|
283
|
+
static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k)
|
284
|
+
{
|
285
|
+
// note I leave an extra slot here at the end because of the way our binary insert works
|
286
|
+
kresult k_list[MAX_K + 1];
|
287
|
+
int k_len = 0;
|
288
|
+
float k_dist = INT_MAX;
|
289
|
+
int ki = NUM2INT(k);
|
290
|
+
VALUE ary;
|
291
|
+
int i;
|
292
|
+
KDTREEP;
|
293
|
+
|
294
|
+
if (ki < 1) {
|
295
|
+
ki = 1;
|
296
|
+
} else if (ki > MAX_K) {
|
297
|
+
ki = MAX_K;
|
298
|
+
}
|
299
|
+
kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0, k_list, &k_len, &k_dist);
|
300
|
+
|
301
|
+
// convert result to ruby array
|
302
|
+
ary = rb_ary_new();
|
303
|
+
for (i = 0; i < k_len; ++i) {
|
304
|
+
rb_ary_push(ary, INT2NUM(kdtreep->nodes[k_list[i].index].id));
|
305
|
+
}
|
306
|
+
return ary;
|
307
|
+
}
|
308
|
+
|
309
|
+
static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth, kresult *k_list, int *k_len, float *k_dist)
|
310
|
+
{
|
311
|
+
struct kdtree_node *n;
|
312
|
+
float ad;
|
313
|
+
int nearer, further;
|
314
|
+
float dx;
|
315
|
+
int lo, hi;
|
316
|
+
|
317
|
+
if (i == -1) {
|
318
|
+
return;
|
319
|
+
}
|
320
|
+
|
321
|
+
n = kdtreep->nodes + i;
|
322
|
+
|
323
|
+
ad = (depth % 2) ? (x - n->x) : (y - n->y);
|
324
|
+
|
325
|
+
//
|
326
|
+
// recurse near, and then perhaps far as well
|
327
|
+
//
|
328
|
+
|
329
|
+
if (ad <= 0) {
|
330
|
+
nearer = n->left; further = n->right;
|
331
|
+
} else {
|
332
|
+
nearer = n->right; further = n->left;
|
333
|
+
}
|
334
|
+
kdtree_nearestk0(kdtreep, nearer, x, y, k, depth + 1, k_list, k_len, k_dist);
|
335
|
+
if (ad * ad < *k_dist) {
|
336
|
+
kdtree_nearestk0(kdtreep, further, x, y, k, depth + 1, k_list, k_len, k_dist);
|
337
|
+
}
|
338
|
+
//
|
339
|
+
// do we beat the old distance?
|
340
|
+
//
|
341
|
+
|
342
|
+
dx = (x - n->x) * (x - n->x);
|
343
|
+
if (dx < *k_dist) {
|
344
|
+
float d = dx + ((y - n->y) * (y - n->y));
|
345
|
+
if (d < *k_dist) {
|
346
|
+
//
|
347
|
+
// find spot to insert
|
348
|
+
//
|
349
|
+
lo = 0, hi = *k_len;
|
350
|
+
while (lo < hi) {
|
351
|
+
int mid = (lo + hi) / 2;
|
352
|
+
if (k_list[mid].distance < d) {
|
353
|
+
lo = mid + 1;
|
354
|
+
} else {
|
355
|
+
hi = mid;
|
356
|
+
}
|
357
|
+
}
|
358
|
+
|
359
|
+
//
|
360
|
+
// insert
|
361
|
+
//
|
362
|
+
|
363
|
+
memmove(k_list + lo + 1, k_list + lo, (*k_len - lo) * sizeof(struct kresult));
|
364
|
+
k_list[lo].index = i;
|
365
|
+
k_list[lo].distance = d;
|
366
|
+
|
367
|
+
//
|
368
|
+
// adjust len/dist if necessary
|
369
|
+
//
|
370
|
+
|
371
|
+
if (*k_len < k) {
|
372
|
+
++(*k_len);
|
373
|
+
} else {
|
374
|
+
*k_dist = k_list[k - 1].distance;
|
375
|
+
}
|
376
|
+
}
|
377
|
+
}
|
378
|
+
}
|
379
|
+
|
380
|
+
/*
|
381
|
+
* call-seq:
|
382
|
+
* kd.persist(io)
|
383
|
+
*
|
384
|
+
* Writes the tree out to <i>io</i> so you can quickly load it later with
|
385
|
+
* Kdtree.new. This avoids the startup cost of initializing a tree. Apart from a
|
386
|
+
* small header, the size of the file is proportional to the number of points,
|
387
|
+
* requiring 20 bytes per point.
|
388
|
+
*
|
389
|
+
* This file is <b>NOT PORTABLE</b> across different architectures due to endian
|
390
|
+
* issues.
|
391
|
+
*
|
392
|
+
* points = []
|
393
|
+
* points << [47.6, -122.3, 1] # Seattle
|
394
|
+
* points << [45.5, -122.8, 2] # Portland
|
395
|
+
* points << [40.7, -74.0, 3] # New York
|
396
|
+
* kd = Kdtree.new(points)
|
397
|
+
*
|
398
|
+
* # persist the tree to disk
|
399
|
+
* File.open("treefile", "w") { |f| kd.persist(f) }
|
400
|
+
*
|
401
|
+
* ...
|
402
|
+
*
|
403
|
+
* # later, read the tree from disk
|
404
|
+
* kd2 = File.open("treefile") { |f| Kdtree.new(f) }
|
405
|
+
*/
|
406
|
+
static VALUE kdtree_persist(VALUE kdtree, VALUE io)
|
407
|
+
{
|
408
|
+
KDTREEP;
|
409
|
+
|
410
|
+
if (!rb_respond_to(io, rb_intern("write"))) {
|
411
|
+
rb_raise(rb_eTypeError, "instance of IO needed");
|
412
|
+
}
|
413
|
+
if (rb_respond_to(io, id_binmode)) {
|
414
|
+
rb_funcall(io, id_binmode, 0);
|
415
|
+
}
|
416
|
+
|
417
|
+
write_all(io, KDTREE_MAGIC, 4);
|
418
|
+
write_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
|
419
|
+
write_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
|
420
|
+
return io;
|
421
|
+
}
|
422
|
+
|
423
|
+
/*
|
424
|
+
* call-seq:
|
425
|
+
* kd.to_s => string
|
426
|
+
*
|
427
|
+
* A string that tells you a bit about the tree.
|
428
|
+
*/
|
429
|
+
static VALUE kdtree_to_s(VALUE kdtree)
|
430
|
+
{
|
431
|
+
char buf[256];
|
432
|
+
KDTREEP;
|
433
|
+
|
434
|
+
sprintf(buf, "#<%s:%p nodes=%d>", rb_obj_classname(kdtree), (void*)kdtree, kdtreep->len);
|
435
|
+
return rb_str_new(buf, strlen(buf));
|
436
|
+
}
|
437
|
+
|
438
|
+
//
|
439
|
+
// io helpers
|
440
|
+
//
|
441
|
+
|
442
|
+
static void read_all(VALUE io, void *buf, int len)
|
443
|
+
{
|
444
|
+
VALUE string = rb_funcall(io, id_read, 1, INT2NUM(len));
|
445
|
+
if (NIL_P(string) || RSTRING_LEN(string) != len) {
|
446
|
+
rb_raise(rb_eEOFError, "end of file reached");
|
447
|
+
}
|
448
|
+
memcpy(buf, RSTRING_PTR(string), len);
|
449
|
+
}
|
450
|
+
|
451
|
+
static void write_all(VALUE io, const void *buf, int len)
|
452
|
+
{
|
453
|
+
rb_funcall(io, id_write, 1, rb_str_new(buf, len));
|
454
|
+
}
|
455
|
+
|
456
|
+
//
|
457
|
+
// entry point
|
458
|
+
//
|
459
|
+
|
460
|
+
/*
|
461
|
+
* Kdtree is an insanely fast data structure for finding the nearest
|
462
|
+
* neighbor(s) to a given point. This implementation only supports 2d
|
463
|
+
* points. Also, it only supports static points - there is no way to edit the
|
464
|
+
* tree after it has been initialized. Kdtree should scale to millions of
|
465
|
+
* points, though it's only been tested with around 1 million.
|
466
|
+
*
|
467
|
+
* Once the tree is constructed, it can be searched with nearest and nearestk.
|
468
|
+
*
|
469
|
+
* To avoid the startup costs associated with creating a new tree, use persist
|
470
|
+
* to write the tree to disk. You can then construct the tree later from that
|
471
|
+
* file.
|
472
|
+
*
|
473
|
+
* points = []
|
474
|
+
* points << [47.6, -122.3, 1] # Seattle
|
475
|
+
* points << [45.5, -122.8, 2] # Portland
|
476
|
+
* points << [40.7, -74.0, 3] # New York
|
477
|
+
* kd = Kdtree.new(points)
|
478
|
+
*
|
479
|
+
* # which city is closest to San Francisco?
|
480
|
+
* kd.nearest(34.1, -118.2) #=> 2
|
481
|
+
* # which two cities are closest to San Francisco?
|
482
|
+
* kd.nearestk(34.1, -118.2, 2) #=> [2, 1]
|
483
|
+
*
|
484
|
+
* For more information on kd trees, see:
|
485
|
+
*
|
486
|
+
* http://en.wikipedia.org/wiki/Kd-tree
|
487
|
+
*/
|
488
|
+
void Init_kdtree()
|
489
|
+
{
|
490
|
+
static VALUE clazz;
|
491
|
+
|
492
|
+
clazz = rb_define_class("Kdtree", rb_cObject);
|
493
|
+
|
494
|
+
rb_define_alloc_func(clazz, kdtree_alloc);
|
495
|
+
rb_define_method(clazz, "initialize", kdtree_initialize, 1);
|
496
|
+
rb_define_method(clazz, "nearest", kdtree_nearest, 2);
|
497
|
+
rb_define_method(clazz, "nearestk", kdtree_nearestk, 3);
|
498
|
+
rb_define_method(clazz, "persist", kdtree_persist, 1);
|
499
|
+
rb_define_method(clazz, "to_s", kdtree_to_s, 0);
|
500
|
+
|
501
|
+
// function ids
|
502
|
+
id_binmode = rb_intern("binmode");
|
503
|
+
id_read = rb_intern("read");
|
504
|
+
id_write = rb_intern("write");
|
505
|
+
}
|
data/kdtree.gemspec
ADDED
@@ -0,0 +1,22 @@
|
|
1
|
+
Gem::Specification.new do |s|
|
2
|
+
s.name = "sc2ai-kdtree"
|
3
|
+
s.version = "0.4.1"
|
4
|
+
|
5
|
+
s.authors = ["Adam Doppelt"]
|
6
|
+
s.email = ["amd@gurge.com"]
|
7
|
+
s.homepage = "http://github.com/gurgeous/kdtree"
|
8
|
+
s.license = "MIT"
|
9
|
+
s.summary = "A fork of gem 'kdtree', which builds on Windows and depends on sc2ai."
|
10
|
+
s.description = <<EOF
|
11
|
+
A fork of gem 'kdtree', which builds on Windows and depends on sc2ai.
|
12
|
+
Use the original 'kdtree' instead. See homepage link.
|
13
|
+
EOF
|
14
|
+
|
15
|
+
s.add_development_dependency "minitest", "~> 5.0"
|
16
|
+
s.add_development_dependency "rake-compiler", "~> 1.0"
|
17
|
+
|
18
|
+
s.files = `git ls-files`.split("\n")
|
19
|
+
s.test_files = `git ls-files -- {test,spec,features}/*`.split("\n")
|
20
|
+
s.extensions = ["ext/kdtree/extconf.rb"]
|
21
|
+
s.require_paths = ["lib"]
|
22
|
+
end
|
data/lib/kdtree.rb
ADDED
@@ -0,0 +1 @@
|
|
1
|
+
require "kdtree.so"
|
data/test/test_kdtree.rb
ADDED
@@ -0,0 +1,148 @@
|
|
1
|
+
require "benchmark"
|
2
|
+
require "kdtree"
|
3
|
+
require "tempfile"
|
4
|
+
require "minitest/autorun"
|
5
|
+
|
6
|
+
#
|
7
|
+
# create a tree
|
8
|
+
#
|
9
|
+
|
10
|
+
class KdtreeTest < Minitest::Test
|
11
|
+
TMP = "#{Dir.tmpdir}/kdtree_test"
|
12
|
+
|
13
|
+
def setup
|
14
|
+
@points = (0...1000).map { |i| [rand_coord, rand_coord, i] }
|
15
|
+
@kdtree = Kdtree.new(@points)
|
16
|
+
end
|
17
|
+
|
18
|
+
def teardown
|
19
|
+
File.unlink(TMP) if File.exist?(TMP)
|
20
|
+
end
|
21
|
+
|
22
|
+
def test_nearest
|
23
|
+
100.times do
|
24
|
+
pt = [rand_coord, rand_coord]
|
25
|
+
|
26
|
+
# kdtree search
|
27
|
+
id = @kdtree.nearest(pt[0], pt[1])
|
28
|
+
kdpt = @points[id]
|
29
|
+
|
30
|
+
# slow search
|
31
|
+
sortpt = @points.sort_by { |i| distance(i, pt) }.first
|
32
|
+
|
33
|
+
# assert
|
34
|
+
kdd = distance(kdpt, pt)
|
35
|
+
sortd = distance(sortpt, pt)
|
36
|
+
assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
|
37
|
+
end
|
38
|
+
end
|
39
|
+
|
40
|
+
def test_nearestk
|
41
|
+
100.times do
|
42
|
+
pt = [rand_coord, rand_coord]
|
43
|
+
|
44
|
+
# kdtree search
|
45
|
+
list = @kdtree.nearestk(pt[0], pt[1], 5)
|
46
|
+
kdpt = @points[list.last]
|
47
|
+
|
48
|
+
# slow search
|
49
|
+
sortpt = @points.sort_by { |i| distance(i, pt) }[list.length - 1]
|
50
|
+
|
51
|
+
# assert
|
52
|
+
kdd = distance(kdpt, pt)
|
53
|
+
sortd = distance(sortpt, pt)
|
54
|
+
assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
|
55
|
+
end
|
56
|
+
end
|
57
|
+
|
58
|
+
def test_persist
|
59
|
+
# write
|
60
|
+
File.open(TMP, "w") { |f| @kdtree.persist(f) }
|
61
|
+
# read
|
62
|
+
kdtree2 = File.open(TMP, "r") { |f| Kdtree.new(f) }
|
63
|
+
|
64
|
+
# now test some random points
|
65
|
+
100.times do
|
66
|
+
pt = [rand_coord, rand_coord]
|
67
|
+
id1 = @kdtree.nearest(*pt)
|
68
|
+
id2 = kdtree2.nearest(*pt)
|
69
|
+
assert(id1 == id2, "kdtree2 differed from kdtree")
|
70
|
+
end
|
71
|
+
end
|
72
|
+
|
73
|
+
def test_bad_magic
|
74
|
+
File.open(TMP, "w") { |f| f.puts "That ain't right" }
|
75
|
+
assert_raises RuntimeError do
|
76
|
+
File.open(TMP, "r") { |f| Kdtree.new(f) }
|
77
|
+
end
|
78
|
+
end
|
79
|
+
|
80
|
+
def test_eof
|
81
|
+
File.open(TMP, "w") { |f| @kdtree.persist(f) }
|
82
|
+
bytes = File.read(TMP)
|
83
|
+
|
84
|
+
[2, 10, 100].each do |len|
|
85
|
+
File.open(TMP, "w") { |f| f.write(bytes[0, len]) }
|
86
|
+
assert_raises EOFError do
|
87
|
+
File.open(TMP, "r") { |f| Kdtree.new(f) }
|
88
|
+
end
|
89
|
+
end
|
90
|
+
end
|
91
|
+
|
92
|
+
def dont_test_speed
|
93
|
+
sizes = [1, 100, 1000, 10000, 100000, 1000000]
|
94
|
+
ks = [1, 5, 50, 255]
|
95
|
+
sizes.each do |s|
|
96
|
+
points = (0...s).map { |i| [rand_coord, rand_coord, i] }
|
97
|
+
|
98
|
+
# build
|
99
|
+
Benchmark.bm(17) do |bm|
|
100
|
+
kdtree = nil
|
101
|
+
bm.report "build" do
|
102
|
+
kdtree = Kdtree.new(points)
|
103
|
+
end
|
104
|
+
bm.report "persist" do
|
105
|
+
File.open(TMP, "w") { |f| kdtree.persist(f) }
|
106
|
+
end
|
107
|
+
bm.report "read" do
|
108
|
+
File.open(TMP, "r") { |f| Kdtree.new(f) }
|
109
|
+
end
|
110
|
+
|
111
|
+
ks.each do |k|
|
112
|
+
bm.report "100 queries (#{k})" do
|
113
|
+
100.times do
|
114
|
+
if k == 1
|
115
|
+
kdtree.nearest(rand_coord, rand_coord)
|
116
|
+
else
|
117
|
+
kdtree.nearestk(rand_coord, rand_coord, k)
|
118
|
+
end
|
119
|
+
end
|
120
|
+
end
|
121
|
+
end
|
122
|
+
end
|
123
|
+
puts
|
124
|
+
end
|
125
|
+
end
|
126
|
+
|
127
|
+
protected
|
128
|
+
|
129
|
+
def distance(a, b)
|
130
|
+
x, y = a[0] - b[0], a[1] - b[1]
|
131
|
+
x * x + y * y
|
132
|
+
end
|
133
|
+
|
134
|
+
def rand_coord
|
135
|
+
rand(0) * 10 - 5
|
136
|
+
end
|
137
|
+
end
|
138
|
+
|
139
|
+
# running dont_test_speed on my i5 2.8ghz:
|
140
|
+
#
|
141
|
+
# user system total real
|
142
|
+
# build 3.350000 0.020000 3.370000 ( 3.520528)
|
143
|
+
# persist 0.150000 0.020000 0.170000 ( 0.301963)
|
144
|
+
# read 0.280000 0.000000 0.280000 ( 0.432676)
|
145
|
+
# 100 queries (1) 0.000000 0.000000 0.000000 ( 0.000319)
|
146
|
+
# 100 queries (5) 0.000000 0.000000 0.000000 ( 0.000412)
|
147
|
+
# 100 queries (50) 0.000000 0.000000 0.000000 ( 0.001417)
|
148
|
+
# 100 queries (255) 0.000000 0.000000 0.000000 ( 0.006268)
|
metadata
ADDED
@@ -0,0 +1,86 @@
|
|
1
|
+
--- !ruby/object:Gem::Specification
|
2
|
+
name: sc2ai-kdtree
|
3
|
+
version: !ruby/object:Gem::Version
|
4
|
+
version: 0.4.1
|
5
|
+
platform: ruby
|
6
|
+
authors:
|
7
|
+
- Adam Doppelt
|
8
|
+
autorequire:
|
9
|
+
bindir: bin
|
10
|
+
cert_chain: []
|
11
|
+
date: 2024-02-06 00:00:00.000000000 Z
|
12
|
+
dependencies:
|
13
|
+
- !ruby/object:Gem::Dependency
|
14
|
+
name: minitest
|
15
|
+
requirement: !ruby/object:Gem::Requirement
|
16
|
+
requirements:
|
17
|
+
- - "~>"
|
18
|
+
- !ruby/object:Gem::Version
|
19
|
+
version: '5.0'
|
20
|
+
type: :development
|
21
|
+
prerelease: false
|
22
|
+
version_requirements: !ruby/object:Gem::Requirement
|
23
|
+
requirements:
|
24
|
+
- - "~>"
|
25
|
+
- !ruby/object:Gem::Version
|
26
|
+
version: '5.0'
|
27
|
+
- !ruby/object:Gem::Dependency
|
28
|
+
name: rake-compiler
|
29
|
+
requirement: !ruby/object:Gem::Requirement
|
30
|
+
requirements:
|
31
|
+
- - "~>"
|
32
|
+
- !ruby/object:Gem::Version
|
33
|
+
version: '1.0'
|
34
|
+
type: :development
|
35
|
+
prerelease: false
|
36
|
+
version_requirements: !ruby/object:Gem::Requirement
|
37
|
+
requirements:
|
38
|
+
- - "~>"
|
39
|
+
- !ruby/object:Gem::Version
|
40
|
+
version: '1.0'
|
41
|
+
description: |
|
42
|
+
A fork of gem 'kdtree', which builds on Windows and depends on sc2ai.
|
43
|
+
Use the original 'kdtree' instead. See homepage link.
|
44
|
+
email:
|
45
|
+
- amd@gurge.com
|
46
|
+
executables: []
|
47
|
+
extensions:
|
48
|
+
- ext/kdtree/extconf.rb
|
49
|
+
extra_rdoc_files: []
|
50
|
+
files:
|
51
|
+
- ".github/workflows/ci.yml"
|
52
|
+
- ".gitignore"
|
53
|
+
- Gemfile
|
54
|
+
- LICENSE
|
55
|
+
- README.md
|
56
|
+
- Rakefile
|
57
|
+
- ext/kdtree/extconf.rb
|
58
|
+
- ext/kdtree/kdtree.c
|
59
|
+
- kdtree.gemspec
|
60
|
+
- lib/kdtree.rb
|
61
|
+
- test/test_kdtree.rb
|
62
|
+
homepage: http://github.com/gurgeous/kdtree
|
63
|
+
licenses:
|
64
|
+
- MIT
|
65
|
+
metadata: {}
|
66
|
+
post_install_message:
|
67
|
+
rdoc_options: []
|
68
|
+
require_paths:
|
69
|
+
- lib
|
70
|
+
required_ruby_version: !ruby/object:Gem::Requirement
|
71
|
+
requirements:
|
72
|
+
- - ">="
|
73
|
+
- !ruby/object:Gem::Version
|
74
|
+
version: '0'
|
75
|
+
required_rubygems_version: !ruby/object:Gem::Requirement
|
76
|
+
requirements:
|
77
|
+
- - ">="
|
78
|
+
- !ruby/object:Gem::Version
|
79
|
+
version: '0'
|
80
|
+
requirements: []
|
81
|
+
rubygems_version: 3.5.3
|
82
|
+
signing_key:
|
83
|
+
specification_version: 4
|
84
|
+
summary: A fork of gem 'kdtree', which builds on Windows and depends on sc2ai.
|
85
|
+
test_files:
|
86
|
+
- test/test_kdtree.rb
|