--- Final Solver (used for performance test) ---
import numpy as np
import numba

@numba.njit
def find_theta(y):
    left = 0
    right = len(y)
    sum_G = 0.0
    count_G = 0
    
    while left < right:
        # Median of 3 pivot
        mid = left + (right - left) // 2
        a = y[left]
        b = y[mid]
        c = y[right - 1]
        
        if a > b:
            if b > c:
                pivot = b
            elif a > c:
                pivot = c
            else:
                pivot = a
        else:
            if a > c:
                pivot = a
            elif b > c:
                pivot = c
            else:
                pivot = b
        
        sum_L = 0.0
        len_L = 0
        for i in range(left, right):
            if y[i] > pivot:
                sum_L += y[i]
                len_L += 1
                
        val = (sum_G + sum_L) - pivot * (count_G + len_L)
        
        if val == 1.0:
            return pivot
        elif val < 1.0:
            new_right = left
            for i in range(left, right):
                if y[i] < pivot:
                    tmp = y[i]
                    y[i] = y[new_right]
                    y[new_right] = tmp
                    new_right += 1
                else:
                    sum_G += y[i]
                    count_G += 1
            right = new_right
        else:
            new_right = left
            for i in range(left, right):
                if y[i] > pivot:
                    tmp = y[i]
                    y[i] = y[new_right]
                    y[new_right] = tmp
                    new_right += 1
            right = new_right
            
    if count_G == 0:
        return 0.0
    return (sum_G - 1.0) / count_G

@numba.njit
def project_fast(y):
    y_copy = y.copy()
    theta = find_theta(y_copy)
    x = np.empty_like(y)
    for i in range(len(y)):
        val = y[i] - theta
        x[i] = val if val > 0 else 0.0
    return x

class Solver:
    def __init__(self):
        # Warmup the Numba JIT compiler so it doesn't count towards runtime
        _ = project_fast(np.array([1.0, 2.0], dtype=np.float64))
        
    def solve(self, problem: dict, **kwargs) -> dict:
        y = np.array(problem.get("y"), dtype=np.float64).flatten()
        x = project_fast(y)
        return {"solution": x}
--- End Solver ---
Running performance test...
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /tests
plugins: anyio-4.13.0, jaxtyping-0.3.9
collected 3 items

../tests/test_outputs.py .
--- Performance Summary ---
Validity: True
Total Baseline Time: 6.0000s
Total Solver Time:   2.5169s
Raw Speedup:         2.3839 x
Final Reward (Score): 2.3839
---------------------------
..

==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
PASSED ../tests/test_outputs.py::test_solver_speedup
======================== 3 passed in 246.80s (0:04:06) =========================
