# BK Tree implementation

I'm working on a system that can search for similar images. This involves being able to search by edit-distance, and as such I've implemented a special data-structure called a BK tree. Some notes here.

The core idea here is to be able to search for items by edit distance, in this case Hamming distance between discrete items. My hashes for this system are stored as 64 bit integers, but fundamentally they can be considered bit-fields, and are treated as such.

In any event, I have some test code that seems to work fairly well, and I'd greatly appreciate some input.

from libc.stdint cimport uint64_t

# Compute number of bits that are not common between a and b.
# return value is a plain integer
cdef uint64_t hamming(uint64_t a, uint64_t b):

cdef uint64_t x
cdef int tot

tot = 0

x = (a ^ b)
while x > 0:
tot += x & 1
x >>= 1

cdef class BkHammingNode(object):

cdef uint64_t nodeHash
cdef set nodeData
cdef dict children

def __init__(self, nodeHash, nodeData):
self.nodeData = set((nodeData, ))
self.children = {}
self.nodeHash = nodeHash

# Insert phash nodeHash into tree, with the associated data nodeData
cpdef insert(self, uint64_t nodeHash, nodeData):

# If the current node has the same has as the data we're inserting,
# add the data to the current node's data set
if nodeHash == self.nodeHash:
return

# otherwise, calculate the edit distance between the new phash and the current node's hash,
# and either recursively insert the data, or create a new child node for the phash
distance = hamming(self.nodeHash, nodeHash)
if not distance in self.children:
self.children[distance] = BkHammingNode(nodeHash, nodeData)
else:
self.children[distance].insert(nodeHash, nodeData)

# Remove node with hash nodeHash and accompanying data nodeData from the tree.
# Returns list of children that must be re-inserted (or false if no children need to be updated),
# number of nodes deleted, and number of nodes that were moved as a 3-tuple.
cpdef remove(self, uint64_t nodeHash, nodeData):
cdef uint64_t deleted = 0
cdef uint64_t moved = 0

# If the node we're on matches the hash we want to delete exactly:
if nodeHash == self.nodeHash:

# Remove the node data associated with the hash we want to remove
self.nodeData.remove(nodeData)

# If we've emptied out the node of data, return all our children so the parent can
# graft the children into the tree in the appropriate place
if not self.nodeData:
# 1 deleted node, 0 moved nodes, return all children for reinsertion by parent
# Parent will pop this node, and reinsert all it's children where apropriate
return list(self), 1, 0

# node has data remaining, do not do any rebuilding
return False, 1, 0

selfDist = hamming(self.nodeHash, nodeHash)

# Removing is basically searching with a distance of zero, and
# then doing operations on the search result.
# As such, scan children where the edit distance between self.nodeHash and the target nodeHash == 0
# Rebuild children where needed
if selfDist in self.children:
moveChildren, childDeleted, childMoved = self.children[selfDist].remove(nodeHash, nodeData)
deleted += childDeleted
moved += childMoved

# If the child returns children, it means the child no longer contains any unique data, so it
# needs to be deleted. As such, pop it from the tree, and re-insert all it's children as
# direct decendents of the current node
if moveChildren:
self.children.pop(selfDist)
for childHash, childData in moveChildren:
self.insert(childHash, childData)
moved += 1

return False, deleted, moved

# Get all child-nodes within an edit distance of distance from baseHash
# returns a set containing the data of each matching node, and a integer representing
# the number of nodes that were touched in the scan.
# Return value is a 2-tuple
cpdef getWithinDistance(self, uint64_t baseHash, int distance):
cdef uint64_t selfDist

selfDist = hamming(self.nodeHash, baseHash)

ret = set()

if selfDist <= distance:
ret |= set(self.nodeData)

touched = 1

for key in self.children.keys():
if key <= selfDist+distance and key >= selfDist-distance:
new, tmpTouch = self.children[key].getWithinDistance(baseHash, distance)
touched += tmpTouch
ret |= new

return ret, touched

def __iter__(self):
for child in self.children.values():
for item in child:
yield item
for item in self.nodeData:
yield (self.nodeHash, item)

class BkHammingTree(object):
root = None

def __init__(self):
self.nodes = 0

def insert(self, nodeHash, nodeData):
if not self.root:
self.root = BkHammingNode(nodeHash, nodeData)
else:
self.root.insert(nodeHash, nodeData)

self.nodes += 1

def remove(self, nodeHash, nodeData):
if not self.root:
raise ValueError("No tree built to remove from!")

rootless, deleted, moved = self.root.remove(nodeHash, nodeData)

# If the node we're deleting is the root node, we need to handle it properly
# if it is, overwrite the root node with one of the values returned, and then
# rebuild the entire tree by reinserting all the nodes
if rootless:
print("Tree root deleted! Rebuilding...")
rootHash, rootData = rootless.pop()
self.root = BkHammingNode(rootHash, rootData)
for childHash, childData in rootless:
self.root.insert(childHash, childData)

self.nodes -= deleted

return deleted, moved

def getWithinDistance(self, baseHash, distance):
if not self.root:
return set()

ret, touched = self.root.getWithinDistance(baseHash, distance)
print("Touched %s tree nodes, or %1.3f%%" % (touched, touched/self.nodes * 100))
print("Discovered %s match(es)" % len(ret))
return ret

def __iter__(self):
for value in self.root:
yield value


This is written in cythonized Python, for performance reasons (it was slow in pure Python). Right now, it's pretty speedy (a 4.1M item tree builds in about 20 seconds), but I have sneaking suspicions I'm missing something, particularly as I still really don't deeply understand metric spaces.

I think the tree is implemented correctly, but it wouldn't surprise me if I'm missing something obvious.

• I know I'm late to the party, but this looks like you'd probably be better off just using PyPy to me. It looks like the kind of stuff PyPy is really good at optimizing. I know it's not always possible to use PyPy, though. – Veedrac Dec 8 '14 at 16:08
• @Veedrac - Interesting. The major issue is that I'm on python 3.4, and need psycopg2. I wanted a chance to play with cython anyways. – Fake Name Dec 8 '14 at 23:43
• It seems you're in luck: pypi.python.org/pypi/psycopg2cffi. Their benchmarks put it in quite a good light, too. Downgrading to 3.2 should be relatively simple if you don't use yield from a ton. – Veedrac Dec 9 '14 at 0:49
• @Veedrac - To be honest, I'm pretty happy performance-wise at this point. I'm mostly interested in correctness, rather then performance optimizations. I have enough other areas where performance is sufficiently terrible that any further effort here would be somewhat moot. – Fake Name Dec 9 '14 at 1:39

I haven't really checked for correctness, but here are some general style points. I haven't taken the time to actually understand the algorithm here; these are surface-level opinions. If you want deeper analysis of the code correctness, I suggest you give me a testing harness (just something to pyximport and run the code), so I get a look at how it's used.

I don't have substantial criticism of the code; it's mostly really neat and readable.

# Compute number of bits that are not common between a and b.
# return value is a plain integer
cdef uint64_t hamming(uint64_t a, uint64_t b):


write

cdef uint64_t hamming(uint64_t a, uint64_t b):
"""
Compute number of bits that are not common between a and b.
return value is a plain integer
"""


If you're actually targetting 3.4 only, don't write (object) in the inheritance list. If you're potentially wanting 2.x compatibility as well, though, it's good to keep it.

You have

    cdef set nodeData
cdef dict children


It's worth noting that these aren't much faster than bog-standard untyped attributes, but as long as you're aware they are a good datatype. This is actually why I expect PyPy will give better speed advantages.

Instead of set((nodeData, )) just write self.nodeData = {nodeData}.

In remove you have cdef uint64_t deleted = 0, but only add to it once. Giving this a type is pointless, especially because your return casts it to a Python type. A similar but lesser concern exists for moved.

Instead of self.children.pop(selfDist) do del self.children[selfDist]. pop is meant for if you use the return value.

You write:

        ret = set()

if selfDist <= distance:
ret |= set(self.nodeData)


It seems to me that

        ret = set()

if selfDist <= distance
ret = set(self.nodeData)


would be better. If not, try ret.update(self.nodeData).

You don't need to call .keys() here:

        for key in self.children.keys():


This:

            if key <= selfDist+distance and key >= selfDist-distance:


should just be

            if selfDist-distance <= key <= selfDist+distance:


or even

            if abs(key - selfDist) <= distance:


I don't know how this would work with selfDist being a uint64_t and key being a Python int, but if you just remove the cdef uint64_t it'll likely work.

You should probably avoid class variables like root = None and just add them to __init__. Class variables are basically global variables, and although having an immutable global as fallback works, it goes counter to how I'd expect the language to be used.

• Regarding the points about set/dict, yeah, I know they're not that great. I was considering trying to use some of the Boost::Set/Boost::Dict classes at one point (cython can interoperate with C++, apparently), but it wound up not being to necessary. I wrote the code with that scaffolding in place for that reason. I may still revisit it for memory efficiency issues at some point. – Fake Name Dec 9 '14 at 2:22
• The global root is intentional, I want to fail noisily if I have multiple instances (there's some threading stuff on top of this that has had it's share of issues). – Fake Name Dec 9 '14 at 2:25
• I somehow didn't know you could initialize sets as {stuff}. Cool! – Fake Name Dec 9 '14 at 2:27
• There's something subtle going on with the if key <= selfDist+distance and key >= selfDist-distance: bit. If I change it to if abs(key - selfDist) <= distance:, my unit tests fail. There's some confusing unsigned/signed conversion stuff going on in there. – Fake Name Dec 9 '14 at 2:32
• The docstrings point is very valid, though. As is the simple overwriting of ret with a new set if needed. – Fake Name Dec 9 '14 at 2:33