qw-gallery-scenery/python/hamming_search.pyx
2021-05-06 11:52:53 +03:00

47 lines
1.7 KiB
Cython

from __future__ import print_function
cimport cython
from libc.stdint cimport uint64_t
from libcpp.vector cimport vector
from libcpp.pair cimport pair
cdef extern int __builtin_popcountll(unsigned long long) nogil
cdef extern from "<algorithm>" namespace "std":
void partial_sort[RandomAccessIterator](RandomAccessIterator first, RandomAccessIterator middle, RandomAccessIterator last) nogil
@cython.boundscheck(False)
@cython.wraparound(False)
cdef vector[pair[ int, int ]] _hamming_knn(uint64_t[:] arr1,uint64_t[:,:] arr2,int[:] image_ids,int n) nogil:
cdef vector[pair[ int, int ]] v=vector[pair[ int, int ]](arr2.shape[0])
cdef int i,j=0
cdef int sum=0
for i in range(arr2.shape[0]):
sum=0
for j in range(4):
sum+=__builtin_popcountll(arr1[j]^arr2[i][j])
v[i]=pair[ int, int ](sum, image_ids[i])
partial_sort[vector[pair[ int, int ]].iterator](v.begin(), v.begin()+n, v.end())
v.resize(n)
return v
@cython.boundscheck(False)
@cython.wraparound(False)
cdef vector[pair[ int, int ]] _hamming_range_query(uint64_t[:] arr1,uint64_t[:,:] arr2,int[:] image_ids,int query_range) nogil:
cdef vector[pair[ int, int ]] v
cdef int i,j=0
cdef int sum=0
for i in range(arr2.shape[0]):
sum=0
for j in range(4):
sum+=__builtin_popcountll(arr1[j]^arr2[i][j])
if sum<=query_range:
v.push_back(pair[ int, int ](sum, image_ids[i]))
return v
def hamming_knn(uint64_t[:] array,uint64_t[:,:] array2,int[:] image_ids,int n):
return _hamming_knn(array,array2,image_ids,n)
def hamming_range_query(uint64_t[:] array,uint64_t[:,:] array2,int[:] image_ids,int query_range):
return _hamming_range_query(array,array2,image_ids,query_range)