stable_partition_test.go

package main

import (
	"slices"
	"strconv"
	"testing"
)

func Test_partition(t *testing.T) {
	t.Run("strings", func(t *testing.T) {
		input := []string{
			"a", "1", "2", "b", "c", "3",
		}
		integers, notIntegers := partitionSlices(input, func(in string) bool {
			_, err := strconv.Atoi(in)
			return err == nil
		})
		if exp := []string{"1", "2", "3"}; !slices.Equal(exp, integers) {
			t.Errorf("expected lower %#v but got %#v", exp, integers)
		}
		if exp := []string{"a", "b", "c"}; !slices.Equal(exp, notIntegers) {
			t.Errorf("expected upper %#v but got %#v", exp, notIntegers)
		}
	})

	for _, tt := range []struct {
		Name  string
		Input []int
		Even  []int
		Odd   []int
	}{
		{
			Name:  "empty",
			Input: []int{},
			Even:  []int{}, Odd: []int{},
		},
		{
			Name:  "one",
			Input: []int{1},
			Even:  []int{}, Odd: []int{1},
		},
		{
			Name:  "two",
			Input: []int{2},
			Even:  []int{2}, Odd: []int{},
		},
		{
			Name:  "unsorted",
			Input: []int{1, 2},
			Even:  []int{2}, Odd: []int{1},
		},
		{
			Name:  "sorted",
			Input: []int{2, 3},
			Even:  []int{2}, Odd: []int{3},
		},
		{
			Name:  "all even",
			Input: []int{2, 4, 6},
			Even:  []int{2, 4, 6}, Odd: []int{},
		},
		{
			Name:  "all odd",
			Input: []int{1, 3, 5},
			Even:  []int{}, Odd: []int{1, 3, 5},
		},
		{
			Name:  "same number of each",
			Input: []int{1, 2, 3, 4, 5, 6},
			Even:  []int{2, 4, 6}, Odd: []int{1, 3, 5},
		},
		{
			Name:  "fewer odd",
			Input: []int{2, 3, 4, 5, 6},
			Even:  []int{2, 4, 6}, Odd: []int{3, 5},
		},
		{
			Name:  "fewer even",
			Input: []int{1, 2, 3, 4, 5},
			Even:  []int{2, 4}, Odd: []int{1, 3, 5},
		},
		{
			Name:  "one odd in the middle",
			Input: []int{2, 2, 3, 2, 2},
			Even:  []int{2, 2, 2, 2}, Odd: []int{3},
		},
		{
			Name:  "one even in the middle",
			Input: []int{1, 1, 2, 1, 1},
			Even:  []int{2}, Odd: []int{1, 1, 1, 1},
		},
	} {
		t.Run(tt.Name, func(t *testing.T) {
			even, odd := partitionSlices(tt.Input, isEven)
			if !slices.Equal(tt.Even, even) {
				t.Errorf("expected even %#v but got %#v", tt.Even, even)
			}
			if !slices.Equal(tt.Odd, odd) {
				t.Errorf("expected odd %#v but got %#v", tt.Odd, odd)
			}
		})
	}

}

func isEven(number int) bool { return number%2 == 0 }

func partitionIndex[T any](list []T, lower func(T) bool) int {
	low := list[:0]
	for i, v := range list {
		if lower(v) {
			low = append(low, v)
			continue
		}
		unsorted := list[i+1:]
		for j, w := range unsorted {
			if lower(w) {
				for k := i + j + 1; k > i; k-- {
					list[k], list[k-1] = list[k-1], list[k]
				}
				low = append(low, w)
				if j == len(unsorted)-1 {
					return len(low)
				}
				break
			}
		}
	}
	return len(low)
}

func partitionSlices[T any](list []T, lower func(T) bool) ([]T, []T) {
	i := partitionIndex(list, lower)
	return slices.Clone(list[:i]), slices.Clone(list[i:])
}