Class: Kdtree
- Inherits:
-
Object
- Object
- Kdtree
- Defined in:
- ext/kdtree/kdtree.c
Instance Method Summary collapse
-
#initialize(arg) ⇒ Object
constructor
Returns a new
Kdtree
. -
#nearest(x, y) ⇒ Object
Finds the point closest to x, y and returns the id for that point.
-
#nearestk(x, y, k) ⇒ Array
Finds the k points closest to x, y.
-
#persist(io) ⇒ Object
Writes the tree out to io so you can quickly load it later with Kdtree.new.
-
#to_s ⇒ String
A string that tells you a bit about the tree.
Constructor Details
#new(points) ⇒ Object #new(io) ⇒ Object
Returns a new Kdtree
. To construct a tree, pass an array of points. Each point should be an array of the form [x, y, id]
, where x and y are floats and id is an integer. The id is arbitrary and will be returned to you whenever you search with nearest or nearestk.
# create a new tree
points = []
points << [47.6, -122.3, 1] # Seattle
points << [40.7, -74.0, 2] # New York
kd = Kdtree.new(points)
Alternately, you can pass in an IO object containing a persisted kdtree. This makes it possible to build the tree in advance, persist it, and start it up quickly later. See persist for more information.
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
# File 'ext/kdtree/kdtree.c', line 98
static VALUE kdtree_initialize(VALUE kdtree, VALUE arg)
{
KDTREEP;
if (TYPE(arg) == T_ARRAY) {
// init from array of pints
VALUE points = arg;
int i;
kdtreep->len = RARRAY_LEN(points);
kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
for (i = 0; i < RARRAY_LEN(points); ++i) {
struct kdtree_node *n = kdtreep->nodes + i;
VALUE ptr = rb_ary_entry(points, i);
VALUE v = rb_check_array_type(ptr);
if (NIL_P(v) || RARRAY_LEN(v) != 3) {
continue;
}
n->x = NUM2DBL(rb_ary_entry(v, 0));
n->y = NUM2DBL(rb_ary_entry(v, 1));
n->id = NUM2INT(rb_ary_entry(v, 2));
}
// now build the tree
kdtreep->root = kdtree_build(kdtreep, 0, kdtreep->len, 0);
} else if (rb_respond_to(arg, rb_intern("read"))) {
VALUE io = arg;
char buf[4];
if (rb_respond_to(io, id_binmode)) {
rb_funcall(io, id_binmode, 0);
}
// check magic
read_all(io, buf, 4);
if (memcmp(KDTREE_MAGIC, buf, 4) != 0) {
rb_raise(rb_eRuntimeError, "wrong magic number in kdtree file");
}
// read start of the struct
read_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
// read the nodes
kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
read_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
} else {
rb_raise(rb_eTypeError, "array or IO required to init Kdtree");
}
return kdtree;
}
|
Instance Method Details
#nearest(x, y) ⇒ Object
Finds the point closest to x, y and returns the id for that point. Returns -1 if the tree is empty.
points = []
points << [47.6, -122.3, 1] # Seattle
points << [40.7, -74.0, 2] # New York
kd = Kdtree.new(points)
# which city is closest to Portland?
kd.nearest(45.5, -122.8) #=> 1
# which city is closest to Boston?
kd.nearest(42.4, -71.1) #=> 2
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
# File 'ext/kdtree/kdtree.c', line 201
static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y)
{
int n_index;
float n_dist;
KDTREEP;
n_index = -1;
n_dist = INT_MAX;
kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0, &n_index, &n_dist);
if (n_index == -1) {
return -1;
}
return INT2NUM((kdtreep->nodes + n_index)->id);
}
|
#nearestk(x, y, k) ⇒ Array
Finds the k points closest to x, y. Returns an array of ids, sorted by distance. Returns an empty array if the tree is empty. Note that k is capped at 255.
points = []
points << [47.6, -122.3, 1] # Seattle
points << [45.5, -122.8, 2] # Portland
points << [40.7, -74.0, 3] # New York
kd = Kdtree.new(points)
# which two cities are closest to San Francisco?
kd.nearest(34.1, -118.2) #=> [2, 1]
283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 |
# File 'ext/kdtree/kdtree.c', line 283
static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k)
{
// note I leave an extra slot here at the end because of the way our binary insert works
kresult k_list[MAX_K + 1];
int k_len = 0;
float k_dist = INT_MAX;
int ki = NUM2INT(k);
VALUE ary;
int i;
KDTREEP;
if (ki < 1) {
ki = 1;
} else if (ki > MAX_K) {
ki = MAX_K;
}
kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0, k_list, &k_len, &k_dist);
// convert result to ruby array
ary = rb_ary_new();
for (i = 0; i < k_len; ++i) {
rb_ary_push(ary, INT2NUM(kdtreep->nodes[k_list[i].index].id));
}
return ary;
}
|
#persist(io) ⇒ Object
Writes the tree out to io so you can quickly load it later with Kdtree.new. This avoids the startup cost of initializing a tree. Apart from a small header, the size of the file is proportional to the number of points, requiring 20 bytes per point.
This file is NOT PORTABLE across different architectures due to endian issues.
points = []
points << [47.6, -122.3, 1] # Seattle
points << [45.5, -122.8, 2] # Portland
points << [40.7, -74.0, 3] # New York
kd = Kdtree.new(points)
# persist the tree to disk
File.open("treefile", "w") { |f| kd.persist(f) }
...
# later, read the tree from disk
kd2 = File.open("treefile") { |f| Kdtree.new(f) }
407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 |
# File 'ext/kdtree/kdtree.c', line 407
static VALUE kdtree_persist(VALUE kdtree, VALUE io)
{
VALUE str;
KDTREEP;
if (!rb_respond_to(io, rb_intern("write"))) {
rb_raise(rb_eTypeError, "instance of IO needed");
}
if (rb_respond_to(io, id_binmode)) {
rb_funcall(io, id_binmode, 0);
}
write_all(io, KDTREE_MAGIC, 4);
write_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
write_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
return io;
}
|
#to_s ⇒ String
A string that tells you a bit about the tree.
431 432 433 434 435 436 437 438 |
# File 'ext/kdtree/kdtree.c', line 431
static VALUE kdtree_to_s(VALUE kdtree)
{
char buf[256];
KDTREEP;
sprintf(buf, "#<%s:%p nodes=%d>", rb_obj_classname(kdtree), (void*)kdtree, kdtreep->len);
return rb_str_new(buf, strlen(buf));
}
|