5. Find Median in Infinite Stream

package infinitequantiles

import (
    "errors"
    "math"
    "sync"
)

var ErrInvalidQuantile = errors.New("p2: quantile must be between 0 and 1")

// P2Estimator вычисляет квантиль потока за строго O(1) памяти
type P2Estimator struct {
    mu sync.RWMutex
    p  float64 // целевой квантиль (0.5 для медианы)

    q  [5]float64 // высоты 5 маркеров
    n  [5]int     // фактические позиции
    np [5]float64 // желаемые позиции
    dn [5]float64 // шаг прироста желаемых позиций

    count int
}

func NewP2Estimator(p float64) (*P2Estimator, error) {
    if p <= 0 || p >= 1 {
        return nil, ErrInvalidQuantile
    }

    est := &P2Estimator{
        p: p,
    }

    est.np[0] = 1
    est.np[1] = 1 + 2*p
    est.np[2] = 1 + 4*p
    est.np[3] = 3 + 2*p
    est.np[4] = 5

    est.dn[0] = 0
    est.dn[1] = p / 2
    est.dn[2] = p
    est.dn[3] = (1 + p) / 2
    est.dn[4] = 1

    return est, nil
}

// Add регистрирует новое число в потоке за время O(1)
func (est *P2Estimator) Add(x float64) {
    est.mu.Lock()
    defer est.mu.Unlock()

    est.count++

    // Инициализация первыми 5 числами
    if est.count <= 5 {
        est.q[est.count-1] = x
        if est.count == 5 {
            est.initMarkers()
        }
        return
    }

    // Шаг 1: Определение интервала вставки
    k := 0
    if x < est.q[0] {
        est.q[0] = x
        k = 0
    } else if x < est.q[1] {
        k = 0
    } else if x < est.q[2] {
        k = 1
    } else if x < est.q[3] {
        k = 2
    } else if x < est.q[4] {
        k = 3
    } else {
        est.q[4] = x
        k = 3
    }

    // Шаг 2: Сдвиг фактических позиций
    for i := k + 1; i < 5; i++ {
        est.n[i]++
    }

    // Шаг 3: Обновление желаемых позиций
    for i := 0; i < 5; i++ {
        est.np[i] += est.dn[i]
    }

    // Шаг 4: Коррекция высот маркеров 1, 2, 3
    for i := 1; i <= 3; i++ {
        d := est.np[i] - float64(est.n[i])
        if (d >= 1 && est.n[i+1]-est.n[i] > 1) || (d <= -1 && est.n[i-1]-est.n[i] < -1) {
            step := int(math.Copysign(1, d))
            qNew := est.parabolic(i, step)

            if est.q[i-1] < qNew && qNew < est.q[i+1] {
                est.q[i] = qNew
            } else {
                est.q[i] = est.linear(i, step)
            }
            est.n[i] += step
        }
    }
}

// Median возвращает приближенную медиану за O(1)
func (est *P2Estimator) Median() float64 {
    est.mu.RLock()
    defer est.mu.RUnlock()

    if est.count == 0 {
        return 0
    }
    if est.count <= 5 {
        return est.q[est.count/2]
    }
    return est.q[2] // Маркер 2 всегда хранит квантиль p (медиану)
}

func (est *P2Estimator) initMarkers() {
    // Сортировка первых 5 чисел на месте
    for i := 0; i < 4; i++ {
        for j := 0; j < 4-i; j++ {
            if est.q[j] > est.q[j+1] {
                est.q[j], est.q[j+1] = est.q[j+1], est.q[j]
            }
        }
    }
    est.n = [5]int{1, 2, 3, 4, 5}
}

func (est *P2Estimator) parabolic(i, d int) float64 {
    step := float64(d)
    n0, n1, n2 := float64(est.n[i-1]), float64(est.n[i]), float64(est.n[i+1])
    q0, q1, q2 := est.q[i-1], est.q[i], est.q[i+1]

    return q1 + (step/(n2-n0))*(
        (n1-n0+step)*(q2-q1)/(n2-n1)+
            (n2-n1-step)*(q1-q0)/(n1-n0))
}

func (est *P2Estimator) linear(i, d int) float64 {
    step := float64(d)
    if d > 0 {
        return est.q[i] + step*(est.q[i+1]-est.q[i])/float64(est.n[i+1]-est.n[i])
    }
    return est.q[i] - step*(est.q[i-1]-est.q[i])/float64(est.n[i-1]-est.n[i])
}