github.com/andy2046/gopie@v0.7.0/pkg/quickselect/quickselect.go (about)

     1  // Package quickselect implements Quickselect.
     2  package quickselect
     3  
     4  // Select selects the kth element.
     5  func Select(data []int, k int) int {
     6  	return do(data, 0, len(data)-1, k-1)
     7  }
     8  
     9  func do(a []int, left, right, k int) int {
    10  	for {
    11  		if right <= left+1 {
    12  			if right == left+1 && a[right] < a[left] {
    13  				swap(a, left, right)
    14  			}
    15  			return a[k]
    16  		}
    17  
    18  		middle := (left + right) >> 1
    19  		swap(a, middle, left+1)
    20  
    21  		if a[left] > a[right] {
    22  			swap(a, left, right)
    23  		}
    24  
    25  		if a[left+1] > a[right] {
    26  			swap(a, left+1, right)
    27  		}
    28  
    29  		if a[left] > a[left+1] {
    30  			swap(a, left, left+1)
    31  		}
    32  
    33  		i, j := left+1, right
    34  		pivot := a[left+1]
    35  
    36  		for {
    37  			for ok := true; ok; ok = a[i] < pivot {
    38  				i++
    39  			}
    40  			for ok := true; ok; ok = a[j] > pivot {
    41  				j--
    42  			}
    43  
    44  			if j < i {
    45  				break
    46  			}
    47  
    48  			swap(a, i, j)
    49  		}
    50  
    51  		a[left+1] = a[j]
    52  		a[j] = pivot
    53  
    54  		if j >= k {
    55  			right = j - 1
    56  		}
    57  
    58  		if j <= k {
    59  			left = i
    60  		}
    61  	}
    62  }
    63  
    64  func swap(a []int, i, j int) {
    65  	a[i], a[j] = a[j], a[i]
    66  }