goimagehash/etcs/utils.go

// Copyright 2017 The goimagehash Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package etcs

import "golang.org/x/exp/constraints"

type Number interface {
	constraints.Integer | constraints.Float
}

// MeanOfPixels function returns a mean of pixels.
func MeanOfPixels[N Number](pixels []N) N {
	var m float64 = 0
	lens := float64(len(pixels))
	if lens == 0 {
		return 0
	}

	for _, p := range pixels {
		m += float64(p)
	}
	// print(m," ", lens, "\n")

	return N(m / lens)
}

// MedianOfPixels function returns a median value of pixels.
// It uses quick selection algorithm.
func MedianOfPixels[N Number](pixels []N) N {
	tmp := make([]N, len(pixels))
	copy(tmp, pixels)
	l := len(tmp)
	pos := l / 2
	v := quickSelectMedian(tmp, 0, l-1, pos)
	return v
}

// MedianOfPixelsFast64 function returns a median value of pixels.
// It uses quick selection algorithm.
func MedianOfPixelsFast64[N Number](pixels []N) N {
	tmp := [64]N{}
	copy(tmp[:], pixels)
	l := len(tmp)
	pos := l / 2
	return quickSelectMedian(tmp[:], 0, l-1, pos)
}

func quickSelectMedian[N Number](sequence []N, low, hi, k int) N {
	if low == hi {
		return sequence[k]
	}

	for low < hi {
		pivot := low/2 + hi/2
		pivotValue := sequence[pivot]
		storeIdx := low
		sequence[pivot], sequence[hi] = sequence[hi], sequence[pivot]
		for i := low; i < hi; i++ {
			if sequence[i] < pivotValue {
				sequence[storeIdx], sequence[i] = sequence[i], sequence[storeIdx]
				storeIdx++
			}
		}
		sequence[hi], sequence[storeIdx] = sequence[storeIdx], sequence[hi]
		if k <= storeIdx {
			hi = storeIdx
		} else {
			low = storeIdx + 1
		}
	}

	if len(sequence)%2 == 0 {
		return sequence[k-1]/2 + sequence[k]/2
	}
	return sequence[k]
}