어떤 키의 집합이 있을 때 거기다 쿼리 키를 인풋하면 집합에서 쿼리키와 유사한 키들을 일정 범위로 찾아내는 BK-Tree를 연구차 파이썬으로 만들어봄.

파이썬 안한지 2,3개월 넘은 데다가 Cpp하다가 만지니까 어 ㅅㅂ 이거 졸라 불안불안한데 왜 안 명시적 타입 지정요 하면서도 어찌어찌는 되게 한듯...

생각컨데 빅오는 Levenshtein distance 구현 땜에 졸라게 클 것 같다. 조트망.


참고로 최적화 해야하는데 그냥 될대로 되라하고 개판인 상태로 올림. 실제 돌아가는 구문은 맨 밑에 달랑 5줄임. PC로 봐야 잘 보임 아마.


from collections import deque


class BKTree:
    def __init__(self, distance_func):
        self._tree = None
        self._distance_func = distance_func

    def add_container(self, container):
        for element in container:
            self.add(element)

    def add(self, node):
        if self._tree is None:
            self._tree = (node, {})
            return

        parent, children = self._tree
        while True:
            distance = self._distance_func(node, parent)
            target = children.get(distance)
            if target is None:
                children[distance] = (node, {})
                break
            else:
                parent, children = target

    def search(self, node, radius):
        if self._tree is None:
            return []

        candidates = deque([self._tree])
        result = []
        while candidates:
            candidt, children = candidates.popleft()
            dist = self._distance_func(node, candidt)
            if dist <= radius:
                result.append((dist, candidt))

            low, high = dist - radius, dist + radius
            candidates.extend(c for d, c in children.items()
                              if low <= d <= high)
        return result


def levenshtein(node, current):
    if node is current:
        return 0
    if len(node) is 0:
        return len(current)
    if len(current) is 0:
        return len(node)

    arr_dic = [i for i in range(0, len(current) + 1)]
    cur_dic = [0 for i in range(0, len(current) + 1)]

    for i in range(0, len(node)):
        cur_dic[0] = i + 1
        for j in range(0, len(current)):
            cost = 0 if node[i] is current[j] else 1
            cur_dic[j + 1] = min([cur_dic[j] + 1, arr_dic[j + 1] + 1, arr_dic[j] + cost])

        for j in range(0, len(arr_dic)):
            arr_dic[j] = cur_dic[j]

    return cur_dic[len(current)]


treeInstance = BKTree(levenshtein)
word_list = ['some', 'same', 'soft', 'mole', 'salmon', 'soda']
treeInstance.add_container(word_list)

print(treeInstance.search('sort', 2))