private List < Double > getTopN ( INDArray vec , int N ) { BasicModelUtils . ArrayComparator comparator = new BasicModelUtils . ArrayComparator ( ) ; PriorityQueue < Double [ ] > queue = new PriorityQueue <> ( vec . rows ( ) , comparator ) ; for ( int j = 0 ; j < vec . length ( ) ; j ++ ) { final Double [ ] pair = new Double [ ] { vec . getDouble ( j ) , ( double ) j } ; if ( queue . size ( ) < N ) { queue . add ( pair ) ; } else { Double [ ] head = queue . peek ( ) ; if ( comparator . compare ( pair , head ) > 0 ) { queue . poll ( ) ; queue . add ( pair ) ; } } } List < Double > lowToHighSimLst = new ArrayList <> ( ) ; while ( ! queue . isEmpty ( ) ) { double ind = queue . poll ( ) [ 1 ] ; lowToHighSimLst . add ( ind ) ; } return Lists . reverse ( lowToHighSimLst ) ; }
Get top N elements