groupon-kdtree 0.1.1

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,2 @@
1
+ .
2
+ bin
data/Gemfile ADDED
@@ -0,0 +1,4 @@
1
+ source "http://github.org"
2
+
3
+ # Specify your gem's dependencies in kdtree.gemspec
4
+ gemspec
data/LICENSE ADDED
@@ -0,0 +1,20 @@
1
+ Copyright (c) 2009 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 ADDED
@@ -0,0 +1,3 @@
1
+ The KDTree gem doesn't exist on GitHub, otherwise we'd have forked it.
2
+
3
+ This gem includes a patch for Ruby 1.9.1.
@@ -0,0 +1,2 @@
1
+ require 'bundler'
2
+ Bundler::GemHelper.install_tasks
@@ -0,0 +1,7 @@
1
+ require 'mkmf'
2
+
3
+ if RUBY_VERSION >= "1.9.0"
4
+ $CFLAGS << " -DUSING_RUBY_19"
5
+ end
6
+
7
+ create_makefile("kdtree")
@@ -0,0 +1,502 @@
1
+ #include "ruby.h"
2
+
3
+ #ifdef USING_RUBY_19
4
+ #include "ruby/io.h"
5
+ #else
6
+ #include "rubyio.h"
7
+ #endif
8
+
9
+ #ifndef HAVE_RB_IO_T
10
+ #define rb_io_t OpenFile
11
+ #endif
12
+
13
+ //
14
+ // interface
15
+ //
16
+
17
+ typedef struct kdtree_data
18
+ {
19
+ int root;
20
+ int len;
21
+ struct kdtree_node *nodes;
22
+ } kdtree_data;
23
+
24
+ typedef struct kdtree_node
25
+ {
26
+ float x, y;
27
+ int id;
28
+ int left;
29
+ int right;
30
+ } kdtree_node;
31
+
32
+ #define KDTREEP \
33
+ struct kdtree_data *kdtreep; \
34
+ Data_Get_Struct(kdtree, struct kdtree_data, kdtreep);
35
+
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
+ // helpers
45
+ static int rb_io_fread_cross_platform(char *buf, int len, rb_io_t *fptr); // ruby1.8 and 1.9 compatible
46
+ static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth);
47
+ static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth);
48
+ static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth);
49
+
50
+ #define KDTREE_MAGIC "KdTr"
51
+
52
+ //
53
+ // implementation
54
+ //
55
+
56
+ static VALUE kdtree_alloc(VALUE klass)
57
+ {
58
+ struct kdtree_data *kdtreep;
59
+ VALUE obj = Data_Make_Struct(klass, struct kdtree_data, 0, kdtree_free, kdtreep);
60
+ kdtreep->root = -1;
61
+ return obj;
62
+ }
63
+
64
+ static void kdtree_free(struct kdtree_data *kdtreep)
65
+ {
66
+ if (kdtreep) {
67
+ free(kdtreep->nodes);
68
+ }
69
+ }
70
+
71
+ static void read_all(rb_io_t *fptr, char *buf, int len)
72
+ {
73
+ while (len > 0) {
74
+ int n = rb_io_fread_cross_platform(buf, len, fptr);
75
+ if (n == 0) {
76
+ rb_eof_error();
77
+ }
78
+ buf += n;
79
+ len -= n;
80
+ }
81
+ }
82
+
83
+ static int rb_io_fread_cross_platform(char *buf, int len, rb_io_t *fptr) {
84
+ #ifdef USING_RUBY_19
85
+ return fread(buf, 1, len, rb_io_stdio_file(fptr));
86
+ #else
87
+ return rb_io_fread(buf, len, fptr->f);
88
+ #endif
89
+ }
90
+ /*
91
+ * call-seq:
92
+ * KDTree.new(points) => kdtree
93
+ * KDTree.new(io) => kdtree
94
+ *
95
+ * Returns a new <code>KDTree</code>. To construct a tree, pass an array of
96
+ * <i>points</i>. Each point should be an array of the form <code>[x, y,
97
+ * id]</code>, where <i>x</i> and <i>y</i> are floats and <i>id</i> is an
98
+ * integer. The <i>id</i> is arbitrary and will be returned to you whenever you
99
+ * search with nearest or nearestk.
100
+ *
101
+ * # create a new tree
102
+ * points = []
103
+ * points << [47.6, -122.3, 1] # Seattle
104
+ * points << [40.7, -74.0, 2] # New York
105
+ * kd = KDTree.new(points)
106
+ *
107
+ * Alternately, you can pass in an <i>IO</i> object containing a persisted
108
+ * kdtree. This makes it possible to build the tree in advance, persist it, and
109
+ * start it up quickly later. See persist for more information.
110
+ */
111
+ static VALUE kdtree_initialize(VALUE kdtree, VALUE arg)
112
+ {
113
+ KDTREEP;
114
+
115
+ if (TYPE(arg) == T_ARRAY) {
116
+ // init from array of pints
117
+ VALUE points = arg;
118
+ kdtreep->len = RARRAY_LEN(points);
119
+ kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
120
+
121
+ int i;
122
+ for (i = 0; i < RARRAY_LEN(points); ++i) {
123
+ struct kdtree_node *n = kdtreep->nodes + i;
124
+
125
+ VALUE ptr = RARRAY_PTR(points)[i];
126
+ VALUE v = rb_check_array_type(ptr);
127
+ if (NIL_P(v) || RARRAY_LEN(v) != 3) {
128
+ continue;
129
+ }
130
+ VALUE *a = RARRAY_PTR(ptr);
131
+ n->x = NUM2DBL(a[0]);
132
+ n->y = NUM2DBL(a[1]);
133
+ n->id = NUM2INT(a[2]);
134
+ }
135
+
136
+ // now build the tree
137
+ kdtreep->root = kdtree_build(kdtreep, 0, kdtreep->len, 0);
138
+ } else if (rb_respond_to(arg, rb_intern("read"))) {
139
+ VALUE io = arg;
140
+ if (rb_respond_to(io, rb_intern("binmode"))) {
141
+ rb_funcall2(io, rb_intern("binmode"), 0, 0);
142
+ }
143
+
144
+ rb_io_taint_check(io);
145
+ rb_io_t *fptr;
146
+ GetOpenFile(io, fptr);
147
+ rb_io_check_readable(fptr);
148
+
149
+ // check magic
150
+ char buf[4];
151
+ read_all(fptr, buf, 4);
152
+ if (memcmp(KDTREE_MAGIC, buf, 4) != 0) {
153
+ rb_raise(rb_eRuntimeError, "wrong magic number in kdtree file");
154
+ }
155
+
156
+ // read start of the struct
157
+ read_all(fptr, (char *)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
158
+ // read the nodes
159
+ kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
160
+ read_all(fptr, (char *)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
161
+ } else {
162
+ rb_raise(rb_eTypeError, "array or IO required to init KDTree");
163
+ }
164
+
165
+ return kdtree;
166
+ }
167
+
168
+ static int comparex(const void *pa, const void *pb)
169
+ {
170
+ float a = ((const struct kdtree_node*)pa)->x;
171
+ float b = ((const struct kdtree_node*)pb)->x;
172
+ return (a < b) ? -1 : ((a > b) ? 1 : 0);
173
+ }
174
+
175
+ static int comparey(const void *pa, const void *pb)
176
+ {
177
+ float a = ((const struct kdtree_node*)pa)->y;
178
+ float b = ((const struct kdtree_node*)pb)->y;
179
+ return (a < b) ? -1 : ((a > b) ? 1 : 0);
180
+ }
181
+
182
+ static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth)
183
+ {
184
+ if (max <= min) {
185
+ return -1;
186
+ }
187
+
188
+ // sort nodes from min to max
189
+ int(*compar)(const void *, const void *) = (depth % 2) ? comparex : comparey;
190
+ qsort(kdtreep->nodes + min, max - min, sizeof(struct kdtree_node), compar);
191
+
192
+ int median = (min + max) / 2;
193
+ struct kdtree_node *m = kdtreep->nodes + median;
194
+ m->left = kdtree_build(kdtreep, min, median, depth + 1);
195
+ m->right = kdtree_build(kdtreep, median + 1, max, depth + 1);
196
+ return median;
197
+ }
198
+
199
+ //
200
+ // nearest
201
+ //
202
+
203
+ static int n_index;
204
+ static float n_dist;
205
+
206
+ /*
207
+ * call-seq:
208
+ * kd.nearest(x, y) => id
209
+ *
210
+ * Finds the point closest to <i>x</i>, <i>y</i> and returns the id for that
211
+ * point. Returns -1 if the tree is empty.
212
+ *
213
+ * points = []
214
+ * points << [47.6, -122.3, 1] # Seattle
215
+ * points << [40.7, -74.0, 2] # New York
216
+ * kd = KDTree.new(points)
217
+ *
218
+ * # which city is closest to Portland?
219
+ * kd.nearest(45.5, -122.8) #=> 1
220
+ * # which city is closest to Boston?
221
+ * kd.nearest(42.4, -71.1) #=> 2
222
+ */
223
+ static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y)
224
+ {
225
+ KDTREEP;
226
+
227
+ n_index = -1;
228
+ n_dist = INT_MAX;
229
+ kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0);
230
+ if (n_index == -1) {
231
+ return -1;
232
+ }
233
+ return INT2NUM((kdtreep->nodes + n_index)->id);
234
+ }
235
+
236
+ static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth)
237
+ {
238
+ if (i == -1) {
239
+ return;
240
+ }
241
+
242
+ struct kdtree_node *n = kdtreep->nodes + i;
243
+
244
+ float ad = (depth % 2) ? (x - n->x) : (y - n->y);
245
+
246
+ //
247
+ // recurse near, and perhaps far as well
248
+ //
249
+
250
+ int near, far;
251
+ if (ad <= 0) {
252
+ near = n->left; far = n->right;
253
+ } else {
254
+ near = n->right; far = n->left;
255
+ }
256
+ kdtree_nearest0(kdtreep, near, x, y, depth + 1);
257
+ if (ad * ad < n_dist) {
258
+ kdtree_nearest0(kdtreep, far, x, y, depth + 1);
259
+ }
260
+
261
+ //
262
+ // do we beat the old distance?
263
+ //
264
+
265
+ float dx = (x - n->x) * (x - n->x);
266
+ if (dx < n_dist) {
267
+ float d = dx + ((y - n->y) * (y - n->y));
268
+ if (d < n_dist) {
269
+ n_index = i;
270
+ n_dist = d;
271
+ }
272
+ }
273
+ }
274
+
275
+ //
276
+ // nearestK
277
+ //
278
+
279
+ #define MAX_K 255
280
+
281
+ typedef struct kresult {
282
+ int index;
283
+ float distance;
284
+ } kresult;
285
+ // note I leave an extra slot here at the end because of the way our binary insert works
286
+ static struct kresult k_list[MAX_K + 1];
287
+ static int k_len;
288
+ static float k_dist;
289
+
290
+ /*
291
+ * call-seq:
292
+ * kd.nearestk(x, y, k) => array
293
+ *
294
+ * Finds the <i>k</i> points closest to <i>x</i>, <i>y</i>. Returns an array of
295
+ * ids, sorted by distance. Returns an empty array if the tree is empty. Note
296
+ * that <i>k</i> is capped at 255.
297
+ *
298
+ * points = []
299
+ * points << [47.6, -122.3, 1] # Seattle
300
+ * points << [45.5, -122.8, 2] # Portland
301
+ * points << [40.7, -74.0, 3] # New York
302
+ * kd = KDTree.new(points)
303
+ *
304
+ * # which two cities are closest to San Francisco?
305
+ * kd.nearest(34.1, -118.2) #=> [2, 1]
306
+ */
307
+ static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k)
308
+ {
309
+ KDTREEP;
310
+
311
+ k_len = 0;
312
+ k_dist = INT_MAX;
313
+
314
+ int ki = NUM2INT(k);
315
+ if (ki < 1) {
316
+ ki = 1;
317
+ } else if (ki > MAX_K) {
318
+ ki = MAX_K;
319
+ }
320
+ kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0);
321
+
322
+ // convert result to ruby array
323
+ VALUE ary = rb_ary_new();
324
+ int i;
325
+ for (i = 0; i < k_len; ++i) {
326
+ rb_ary_push(ary, INT2NUM(kdtreep->nodes[k_list[i].index].id));
327
+ }
328
+ return ary;
329
+ }
330
+
331
+ static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth)
332
+ {
333
+ if (i == -1) {
334
+ return;
335
+ }
336
+
337
+ struct kdtree_node *n = kdtreep->nodes + i;
338
+
339
+ float ad = (depth % 2) ? (x - n->x) : (y - n->y);
340
+
341
+ //
342
+ // recurse near, and then perhaps far as well
343
+ //
344
+
345
+ int near, far;
346
+ if (ad <= 0) {
347
+ near = n->left; far = n->right;
348
+ } else {
349
+ near = n->right; far = n->left;
350
+ }
351
+ kdtree_nearestk0(kdtreep, near, x, y, k, depth + 1);
352
+ if (ad * ad < k_dist) {
353
+ kdtree_nearestk0(kdtreep, far, x, y, k, depth + 1);
354
+ }
355
+
356
+ //
357
+ // do we beat the old distance?
358
+ //
359
+
360
+ float dx = (x - n->x) * (x - n->x);
361
+ if (dx < k_dist) {
362
+ float d = dx + ((y - n->y) * (y - n->y));
363
+ if (d < k_dist) {
364
+ //
365
+ // find spot to insert
366
+ //
367
+ int lo = 0, hi = k_len;
368
+ while (lo < hi) {
369
+ int mid = (lo + hi) / 2;
370
+ if (k_list[mid].distance < d) {
371
+ lo = mid + 1;
372
+ } else {
373
+ hi = mid;
374
+ }
375
+ }
376
+
377
+ //
378
+ // insert
379
+ //
380
+
381
+ memmove(k_list + lo + 1, k_list + lo, (k_len - lo) * sizeof(struct kresult));
382
+ k_list[lo].index = i;
383
+ k_list[lo].distance = d;
384
+
385
+ //
386
+ // adjust len/dist if necessary
387
+ //
388
+
389
+ if (k_len < k) {
390
+ ++k_len;
391
+ } else {
392
+ k_dist = k_list[k - 1].distance;
393
+ }
394
+ }
395
+ }
396
+ }
397
+
398
+ /*
399
+ * call-seq:
400
+ * kd.persist(io)
401
+ *
402
+ * Writes the tree out to <i>io</i> so you can quickly load it later with
403
+ * KDTree.new. This avoids the startup cost of initializing a tree. Apart from a
404
+ * small header, the size of the file is proportional to the number of points,
405
+ * requiring 20 bytes per point.
406
+ *
407
+ * This file is <b>NOT PORTABLE</b> across different architectures due to endian
408
+ * issues.
409
+ *
410
+ * points = []
411
+ * points << [47.6, -122.3, 1] # Seattle
412
+ * points << [45.5, -122.8, 2] # Portland
413
+ * points << [40.7, -74.0, 3] # New York
414
+ * kd = KDTree.new(points)
415
+ *
416
+ * # persist the tree to disk
417
+ * File.open("treefile", "w") { |f| kd.persist(f) }
418
+ *
419
+ * ...
420
+ *
421
+ * # later, read the tree from disk
422
+ * kd2 = File.open("treefile") { |f| KDTree.new(f) }
423
+ */
424
+ static VALUE kdtree_persist(VALUE kdtree, VALUE io)
425
+ {
426
+ KDTREEP;
427
+
428
+ if (!rb_respond_to(io, rb_intern("write"))) {
429
+ rb_raise(rb_eTypeError, "instance of IO needed");
430
+ }
431
+ if (rb_respond_to(io, rb_intern("binmode"))) {
432
+ rb_funcall2(io, rb_intern("binmode"), 0, 0);
433
+ }
434
+
435
+ VALUE str = rb_str_buf_new(0);
436
+ rb_str_buf_cat(str, KDTREE_MAGIC, 4);
437
+ rb_str_buf_cat(str, (char*)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
438
+ rb_str_buf_cat(str, (char*)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
439
+ rb_io_write(io, str);
440
+ return io;
441
+ }
442
+
443
+ /*
444
+ * call-seq:
445
+ * kd.to_s => string
446
+ *
447
+ * A string that tells you a bit about the tree.
448
+ */
449
+ static VALUE kdtree_to_s(VALUE kdtree)
450
+ {
451
+ KDTREEP;
452
+
453
+ char buf[256];
454
+ sprintf(buf, "#<%s:%p nodes=%d>", rb_obj_classname(kdtree), (void*)kdtree, kdtreep->len);
455
+ return rb_str_new(buf, strlen(buf));
456
+ }
457
+
458
+ //
459
+ // entry point
460
+ //
461
+
462
+ /*
463
+ * KDTree is an insanely fast data structure for finding the nearest
464
+ * neighbor(s) to a given point. This implementation only supports 2d
465
+ * points. Also, it only supports static points - there is no way to edit the
466
+ * tree after it has been initialized. KDTree should scale to millions of
467
+ * points, though it's only been tested with around 1 million.
468
+ *
469
+ * Once the tree is constructed, it can be searched with nearest and nearestk.
470
+ *
471
+ * To avoid the startup costs associated with creating a new tree, use persist
472
+ * to write the tree to disk. You can then construct the tree later from that
473
+ * file.
474
+ *
475
+ * points = []
476
+ * points << [47.6, -122.3, 1] # Seattle
477
+ * points << [45.5, -122.8, 2] # Portland
478
+ * points << [40.7, -74.0, 3] # New York
479
+ * kd = KDTree.new(points)
480
+ *
481
+ * # which city is closest to San Francisco?
482
+ * kd.nearest(34.1, -118.2) #=> 2
483
+ * # which two cities are closest to San Francisco?
484
+ * kd.nearest(34.1, -118.2) #=> [2, 1]
485
+ *
486
+ * For more information on kd trees, see:
487
+ *
488
+ * http://en.wikipedia.org/wiki/Kd-tree
489
+ */
490
+ void Init_kdtree()
491
+ {
492
+ static VALUE clazz;
493
+
494
+ clazz = rb_define_class("KDTree", rb_cObject);
495
+
496
+ rb_define_alloc_func(clazz, kdtree_alloc);
497
+ rb_define_method(clazz, "initialize", kdtree_initialize, 1);
498
+ rb_define_method(clazz, "nearest", kdtree_nearest, 2);
499
+ rb_define_method(clazz, "nearestk", kdtree_nearestk, 3);
500
+ rb_define_method(clazz, "persist", kdtree_persist, 1);
501
+ rb_define_method(clazz, "to_s", kdtree_to_s, 0);
502
+ }
@@ -0,0 +1,23 @@
1
+ # -*- encoding: utf-8 -*-
2
+ #$:.push File.expand_path("../lib", __FILE__)
3
+ require 'date'
4
+ Gem::Specification.new do |s|
5
+ s.name = "groupon-kdtree"
6
+ s.version = "0.1.1"
7
+ s.platform = Gem::Platform::RUBY
8
+ s.author = "Adam Doppelt"
9
+ s.email = "amd@gurge.com"
10
+ s.homepage = ""
11
+ s.description = %q{KDTree updated for Ruby 1.9.x}
12
+
13
+ s.files = `git ls-files`.split("\n")
14
+ s.test_files = `git ls-files -- {test,spec,features}/*`.split("\n")
15
+ s.require_paths = ["."]
16
+ s.rdoc_options = %w(--exclude test --exclude extconf)
17
+ s.summary = "Blazingly fast 2d kdtree."
18
+ require 'date'
19
+ s.date = Date.today.to_s
20
+ s.required_ruby_version = ">= 1.8.5"
21
+ s.has_rdoc = true
22
+ s.extensions << "ext/extconf.rb"
23
+ end
@@ -0,0 +1,138 @@
1
+ require "#{File.expand_path(File.dirname(__FILE__))}/../ext/kdtree.o"
2
+ require "test/unit"
3
+ require "tempfile"
4
+
5
+ #
6
+ # create a tree
7
+ #
8
+
9
+ class KDTreeTest < Test::Unit::TestCase
10
+ TMP = "#{Dir.tmpdir}/kdtree_test"
11
+
12
+ def test_nearest
13
+ setup_tree(1000)
14
+ 100.times do
15
+ pt = [rand_coord, rand_coord]
16
+
17
+ # kdtree search
18
+ id = @kdtree.nearest(pt[0], pt[1])
19
+ kdpt = @points[id]
20
+
21
+ # slow search
22
+ sortpt = @points.sort_by { |i| distance(i, pt) }.first
23
+
24
+ # assert
25
+ kdd = distance(kdpt, pt)
26
+ sortd = distance(sortpt, pt)
27
+ assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
28
+ end
29
+ end
30
+
31
+ def test_nearestk
32
+ setup_tree(1000)
33
+ 100.times do
34
+ pt = [rand_coord, rand_coord]
35
+
36
+ # kdtree search
37
+ list = @kdtree.nearestk(pt[0], pt[1], 5)
38
+ kdpt = @points[list.last]
39
+
40
+ # slow search
41
+ sortpt = @points.sort_by { |i| distance(i, pt) }[list.length - 1]
42
+
43
+ # assert
44
+ kdd = distance(kdpt, pt)
45
+ sortd = distance(sortpt, pt)
46
+ assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
47
+ end
48
+ end
49
+
50
+ def test_persist
51
+ setup_tree(1000)
52
+
53
+ begin
54
+ # write
55
+ File.open(TMP, "w") { |f| @kdtree.persist(f) }
56
+ # read
57
+ kdtree2 = File.open(TMP, "r") { |f| KDTree.new(f) }
58
+
59
+ # now test some random points
60
+ 100.times do
61
+ pt = [rand_coord, rand_coord]
62
+ id1 = @kdtree.nearest(*pt)
63
+ id2 = kdtree2.nearest(*pt)
64
+ assert(id1 == id2, "kdtree2 differed from kdtree")
65
+ end
66
+ ensure
67
+ File.unlink(TMP)
68
+ end
69
+
70
+ # now test magic problems
71
+ begin
72
+ File.open(TMP, "w") { |f| f.puts "That ain't right" }
73
+ assert_raise RuntimeError do
74
+ File.open(TMP, "r") { |f| KDTree.new(f) }
75
+ end
76
+ ensure
77
+ File.unlink(TMP)
78
+ end
79
+ end
80
+
81
+ def dont_test_speed
82
+ printf("\n")
83
+ sizes = [1, 100, 1000, 10000, 100000, 1000000]
84
+ ks = [1, 5, 50, 255]
85
+ sizes.each do |s|
86
+ points = (0...s).map { |i| [rand_coord, rand_coord, i] }
87
+
88
+ # build
89
+ tm = Time.now
90
+ kdtree = KDTree.new(points)
91
+ printf "build %d took %.6fs\n", s, Time.now - tm
92
+
93
+ begin
94
+ # write
95
+ tm = Time.now
96
+ File.open(TMP, "w") { |f| kdtree.persist(f) }
97
+ printf "write %d took %.6fs\n", s, Time.now - tm
98
+ # read
99
+ tm = Time.now
100
+ File.open(TMP, "r") { |f| KDTree.new(f) }
101
+ printf "read %d took %.6fs\n", s, Time.now - tm
102
+ ensure
103
+ File.unlink(TMP)
104
+ end
105
+
106
+ ks.each do |k|
107
+ total = count = 0
108
+ 100.times do
109
+ tm = Time.now
110
+ if k == 1
111
+ kdtree.nearest(rand_coord, rand_coord)
112
+ else
113
+ kdtree.nearestk(rand_coord, rand_coord, k)
114
+ end
115
+ total += Time.now - tm
116
+ count += 1
117
+ end
118
+ printf "avg query time = %.6fs [%d/%d]\n", total / count, s, k
119
+ end
120
+ end
121
+ end
122
+
123
+ protected
124
+
125
+ def setup_tree(len)
126
+ @points = (0...len).map { |i| [rand_coord, rand_coord, i] }
127
+ @kdtree = KDTree.new(@points)
128
+ end
129
+
130
+ def distance(a, b)
131
+ x, y = a[0] - b[0], a[1] - b[1]
132
+ x * x + y * y
133
+ end
134
+
135
+ def rand_coord
136
+ rand(0) * 10 - 5
137
+ end
138
+ end
metadata ADDED
@@ -0,0 +1,75 @@
1
+ --- !ruby/object:Gem::Specification
2
+ name: groupon-kdtree
3
+ version: !ruby/object:Gem::Version
4
+ prerelease: false
5
+ segments:
6
+ - 0
7
+ - 1
8
+ - 1
9
+ version: 0.1.1
10
+ platform: ruby
11
+ authors:
12
+ - Adam Doppelt
13
+ autorequire:
14
+ bindir: bin
15
+ cert_chain: []
16
+
17
+ date: 2011-08-04 00:00:00 -05:00
18
+ default_executable:
19
+ dependencies: []
20
+
21
+ description: KDTree updated for Ruby 1.9.x
22
+ email: amd@gurge.com
23
+ executables: []
24
+
25
+ extensions:
26
+ - ext/extconf.rb
27
+ extra_rdoc_files: []
28
+
29
+ files:
30
+ - .require_paths
31
+ - Gemfile
32
+ - LICENSE
33
+ - README
34
+ - Rakefile
35
+ - ext/extconf.rb
36
+ - ext/kdtree.c
37
+ - kdtree.gemspec
38
+ - test/test.rb
39
+ has_rdoc: true
40
+ homepage: ""
41
+ licenses: []
42
+
43
+ post_install_message:
44
+ rdoc_options:
45
+ - --exclude
46
+ - test
47
+ - --exclude
48
+ - extconf
49
+ require_paths:
50
+ - .
51
+ required_ruby_version: !ruby/object:Gem::Requirement
52
+ requirements:
53
+ - - ">="
54
+ - !ruby/object:Gem::Version
55
+ segments:
56
+ - 1
57
+ - 8
58
+ - 5
59
+ version: 1.8.5
60
+ required_rubygems_version: !ruby/object:Gem::Requirement
61
+ requirements:
62
+ - - ">="
63
+ - !ruby/object:Gem::Version
64
+ segments:
65
+ - 0
66
+ version: "0"
67
+ requirements: []
68
+
69
+ rubyforge_project:
70
+ rubygems_version: 1.3.6
71
+ signing_key:
72
+ specification_version: 3
73
+ summary: Blazingly fast 2d kdtree.
74
+ test_files:
75
+ - test/test.rb