milvus-io /
milvus
| @@ 599-613 (lines=15) @@ | ||
| 596 | return D, I |
|
| 597 | ||
| 598 | ||
| 599 | def kmax(array, k): |
|
| 600 | """return k largest values (and their indices) of the lines of a |
|
| 601 | float32 array""" |
|
| 602 | m, n = array.shape |
|
| 603 | I = np.zeros((m, k), dtype='int64') |
|
| 604 | D = np.zeros((m, k), dtype='float32') |
|
| 605 | ha = float_minheap_array_t() |
|
| 606 | ha.ids = swig_ptr(I) |
|
| 607 | ha.val = swig_ptr(D) |
|
| 608 | ha.nh = m |
|
| 609 | ha.k = k |
|
| 610 | ha.heapify() |
|
| 611 | ha.addn(n, swig_ptr(array)) |
|
| 612 | ha.reorder() |
|
| 613 | return D, I |
|
| 614 | ||
| 615 | ||
| 616 | def pairwise_distances(xq, xb, mt=METRIC_L2, metric_arg=0): |
|
| @@ 582-596 (lines=15) @@ | ||
| 579 | # Wrapper for a few functions |
|
| 580 | ########################################### |
|
| 581 | ||
| 582 | def kmin(array, k): |
|
| 583 | """return k smallest values (and their indices) of the lines of a |
|
| 584 | float32 array""" |
|
| 585 | m, n = array.shape |
|
| 586 | I = np.zeros((m, k), dtype='int64') |
|
| 587 | D = np.zeros((m, k), dtype='float32') |
|
| 588 | ha = float_maxheap_array_t() |
|
| 589 | ha.ids = swig_ptr(I) |
|
| 590 | ha.val = swig_ptr(D) |
|
| 591 | ha.nh = m |
|
| 592 | ha.k = k |
|
| 593 | ha.heapify() |
|
| 594 | ha.addn(n, swig_ptr(array)) |
|
| 595 | ha.reorder() |
|
| 596 | return D, I |
|
| 597 | ||
| 598 | ||
| 599 | def kmax(array, k): |
|
