github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/sort/kofn.go (about)

     1  package sort
     2  
     3  import "reflect"
     4  
     5  // KSmallest for k smallest
     6  // mutates ns
     7  // not sorted
     8  func KSmallest(slice interface{}, k int, cmp func(j, k int) int) interface{} {
     9  
    10  	v := reflect.ValueOf(slice)
    11  	i := v.Len() / 2
    12  
    13  	pos := PartitionLT(slice, i, cmp) + 1
    14  	if pos == k {
    15  		return v.Slice(0, k).Interface()
    16  	}
    17  
    18  	if pos < k {
    19  		return reflect.AppendSlice(
    20  			v.Slice(0, pos),
    21  			reflect.ValueOf(KLargest(v.Slice(pos+1, v.Len()).Interface(), k-pos, cmp)),
    22  		).Interface()
    23  	}
    24  
    25  	// pos > k
    26  
    27  	return KLargest(v.Slice(0, pos-1).Interface(), k, cmp)
    28  }
    29  
    30  // KLargest for k largest
    31  // mutates ns
    32  func KLargest(slice interface{}, k int, cmp func(j, k int) int) interface{} {
    33  	v := reflect.ValueOf(slice)
    34  	i := v.Len() / 2
    35  
    36  	pos := PartitionGT(slice, i, cmp) + 1
    37  	if pos == k {
    38  		return v.Slice(0, k).Interface()
    39  	}
    40  
    41  	if pos < k {
    42  		return reflect.AppendSlice(
    43  			v.Slice(0, pos),
    44  			reflect.ValueOf(KLargest(v.Slice(pos+1, v.Len()).Interface(), k-pos, cmp)),
    45  		).Interface()
    46  	}
    47  
    48  	// pos > k
    49  
    50  	return KLargest(v.Slice(0, pos-1).Interface(), k, cmp)
    51  }
    52  
    53  // PartitionLT partitions array by i-th element
    54  // mutates ns so that all values less than i-th element are on the left
    55  // assume values are distinct
    56  // returns the pos of i-th element
    57  func PartitionLT(slice interface{}, i int, cmp func(j, k int) int) (pos int) {
    58  	v := reflect.ValueOf(slice)
    59  	swp := reflect.Swapper(slice)
    60  
    61  	size := v.Len()
    62  	for j := 0; j < size; j++ {
    63  		if cmp(j, i) < 0 {
    64  			pos++
    65  		}
    66  	}
    67  
    68  	if i != pos {
    69  		swp(i, pos)
    70  	}
    71  
    72  	ri := pos + 1
    73  	if ri == size {
    74  		return
    75  	}
    76  	for li := 0; li < pos; li++ {
    77  		if cmp(li, pos) > 0 {
    78  			for {
    79  				if cmp(ri, pos) < 0 {
    80  					swp(li, ri)
    81  					ri++
    82  					break
    83  				} else {
    84  					ri++
    85  				}
    86  			}
    87  			if ri == size {
    88  				return
    89  			}
    90  		}
    91  
    92  	}
    93  
    94  	return
    95  }
    96  
    97  // PartitionGT places larger ones on the left
    98  func PartitionGT(slice interface{}, i int, cmp func(j, k int) int) (pos int) {
    99  	v := reflect.ValueOf(slice)
   100  	swp := reflect.Swapper(slice)
   101  
   102  	size := v.Len()
   103  	for j := 0; j < size; j++ {
   104  		if cmp(j, i) > 0 {
   105  			pos++
   106  		}
   107  	}
   108  
   109  	if i != pos {
   110  		swp(i, pos)
   111  	}
   112  
   113  	ri := pos + 1
   114  	if ri == size {
   115  		return
   116  	}
   117  	for li := 0; li < pos; li++ {
   118  		if cmp(li, pos) < 0 {
   119  			for {
   120  				if cmp(ri, pos) > 0 {
   121  					swp(li, ri)
   122  					ri++
   123  					break
   124  				} else {
   125  					ri++
   126  				}
   127  			}
   128  			if ri == size {
   129  				return
   130  			}
   131  		}
   132  
   133  	}
   134  
   135  	return
   136  }