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

     1  package sort
     2  
     3  import (
     4  	"testing"
     5  
     6  	"gotest.tools/assert"
     7  )
     8  
     9  func TestPartition(t *testing.T) {
    10  
    11  	{
    12  		data := [][]int{
    13  			[]int{5, 4, 3, 2, 1},
    14  			[]int{1, 2, 3, 4, 5},
    15  			[]int{1, 3, 2, 5, 4},
    16  		}
    17  
    18  		for _, ns := range data {
    19  			i := 3
    20  			v := ns[i]
    21  			cmp := func(i, j int) int {
    22  				return ns[i] - ns[j]
    23  			}
    24  			pos := PartitionLT(ns, i, cmp)
    25  			verifyLT(t, ns, v, pos)
    26  
    27  		}
    28  	}
    29  
    30  	{
    31  		data := [][]int{
    32  			[]int{5, 4, 3, 2, 1},
    33  			[]int{1, 2, 3, 4, 5},
    34  			[]int{1, 3, 2, 5, 4},
    35  		}
    36  
    37  		for _, ns := range data {
    38  			i := 3
    39  			v := ns[i]
    40  			cmp := func(i, j int) int {
    41  				return ns[i] - ns[j]
    42  			}
    43  			pos := PartitionGT(ns, i, cmp)
    44  			verifyGT(t, ns, v, pos)
    45  		}
    46  	}
    47  
    48  	{
    49  		data := [][]int{
    50  			[]int{5, 4, 3, 2, 1},
    51  			[]int{1, 2, 3, 4, 5},
    52  			[]int{1, 3, 2, 5, 4},
    53  		}
    54  
    55  		for _, ns := range data {
    56  			cmp := func(i, j int) int {
    57  				return ns[i] - ns[j]
    58  			}
    59  			ks := KSmallest(ns, 3, cmp)
    60  			verifyKS(t, ks.([]int), ns, 3)
    61  		}
    62  	}
    63  
    64  	{
    65  		data := [][]int{
    66  			[]int{5, 4, 3, 2, 1},
    67  			[]int{1, 2, 3, 4, 5},
    68  			[]int{1, 3, 2, 5, 4},
    69  		}
    70  
    71  		for _, ns := range data {
    72  			cmp := func(i, j int) int {
    73  				return ns[i] - ns[j]
    74  			}
    75  			kl := KLargest(ns, 3, cmp)
    76  			verifyKL(t, kl.([]int), ns, 3)
    77  		}
    78  	}
    79  
    80  }
    81  
    82  func verifyKL(t *testing.T, ks, ns []int, k int) {
    83  	assert.Assert(t, len(ks) == k, "len(ks) != k")
    84  	m := make(map[int]bool)
    85  	for _, v := range ns {
    86  		m[v] = true
    87  	}
    88  
    89  	for _, v := range ks {
    90  		delete(m, v)
    91  	}
    92  
    93  	for _, v1 := range ks {
    94  		for v2 := range m {
    95  			assert.Assert(t, v2 < v1, "v2>v1")
    96  		}
    97  	}
    98  }
    99  
   100  func verifyKS(t *testing.T, ks, ns []int, k int) {
   101  	assert.Assert(t, len(ks) == k, "len(ks) != k")
   102  	m := make(map[int]bool)
   103  	for _, v := range ns {
   104  		m[v] = true
   105  	}
   106  
   107  	for _, v := range ks {
   108  		delete(m, v)
   109  	}
   110  
   111  	for _, v1 := range ks {
   112  		for v2 := range m {
   113  			assert.Assert(t, v2 > v1, "v2<v1")
   114  		}
   115  	}
   116  }
   117  
   118  func verifyLT(t *testing.T, ns []int, v int, pos int) {
   119  	assert.Assert(t, v == ns[pos], "v != ns[pos]")
   120  	for i := 0; i < len(ns); i++ {
   121  		if i < pos {
   122  			assert.Assert(t, ns[i] < v, "left element not smaller")
   123  		} else if i > pos {
   124  			assert.Assert(t, ns[i] > v, "right element not larger")
   125  		}
   126  	}
   127  }
   128  
   129  func verifyGT(t *testing.T, ns []int, v int, pos int) {
   130  	assert.Assert(t, v == ns[pos], "v != ns[pos]")
   131  	for i := 0; i < len(ns); i++ {
   132  		if i < pos {
   133  			assert.Assert(t, ns[i] > v, "left element not larger")
   134  		} else if i > pos {
   135  			assert.Assert(t, ns[i] < v, "right element not smaller")
   136  		}
   137  	}
   138  }