beam: next greater element variants with brute-force tests

wizard-work/beam/next_greater.py

@@ -0,0 +1,90 @@
+"""Next Greater Element variants.
+
+For each index i, compute the index/value of the next element strictly
+greater than a[i] in the specified direction/order. Returns -1 (or None
+index) when none exists. All variants O(n) via monotonic decreasing
+stack of indices.
+"""
+
+from typing import List, Optional
+
+
+def next_greater_right(nums: List[int]) -> List[int]:
+    """For each i, index of first j>i with nums[j] > nums[i], else -1."""
+    n = len(nums)
+    res = [-1] * n
+    stack: List[int] = []  # indices, values strictly decreasing
+    for i in range(n):
+        while stack and nums[stack[-1]] < nums[i]:
+            res[stack.pop()] = i
+        stack.append(i)
+    return res
+
+
+def next_greater_left(nums: List[int]) -> List[int]:
+    """For each i, index of last j<i with nums[j] > nums[i], else -1."""
+    n = len(nums)
+    res = [-1] * n
+    stack: List[int] = []
+    for i in range(n):
+        while stack and nums[stack[-1]] <= nums[i]:
+            stack.pop()
+        res[i] = stack[-1] if stack else -1
+        stack.append(i)
+    return res
+
+
+def next_greater_circular(nums: List[int]) -> List[int]:
+    """Circular: search wraps around; -1 if none greater in entire array."""
+    n = len(nums)
+    if n == 0:
+        return []
+    res = [-1] * n
+    stack: List[int] = []
+    # traverse 2n-1; only push first pass indices
+    for k in range(2 * n - 1):
+        i = k % n
+        while stack and nums[stack[-1]] < nums[i]:
+            res[stack.pop()] = i
+        if k < n:
+            stack.append(i)
+    return res
+
+
+def values_from_indices(nums: List[int], idxs: List[int]) -> List[Optional[int]]:
+    return [nums[i] if i != -1 else None for i in idxs]
+
+
+# brute-force oracles
+def _brute_right(nums):
+    n = len(nums)
+    out = [-1] * n
+    for i in range(n):
+        for j in range(i + 1, n):
+            if nums[j] > nums[i]:
+                out[i] = j
+                break
+    return out
+
+
+def _brute_left(nums):
+    n = len(nums)
+    out = [-1] * n
+    for i in range(n):
+        for j in range(i - 1, -1, -1):
+            if nums[j] > nums[i]:
+                out[i] = j
+                break
+    return out
+
+
+def _brute_circ(nums):
+    n = len(nums)
+    out = [-1] * n
+    for i in range(n):
+        for step in range(1, n):
+            j = (i + step) % n
+            if nums[j] > nums[i]:
+                out[i] = j
+                break
+    return out

wizard-work/beam/test_next_greater.py

@@ -0,0 +1,39 @@
+import random
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from next_greater import (
+    next_greater_right, next_greater_left, next_greater_circular,
+    _brute_right, _brute_left, _brute_circ,
+)
+
+
+def test_right_random():
+    rng = random.Random(42)
+    for _ in range(2000):
+        n = rng.randint(0, 40)
+        nums = [rng.randint(-5, 5) for _ in range(n)]
+        assert next_greater_right(nums) == _brute_right(nums), nums
+
+
+def test_left_random():
+    rng = random.Random(7)
+    for _ in range(2000):
+        n = rng.randint(0, 40)
+        nums = [rng.randint(-5, 5) for _ in range(n)]
+        assert next_greater_left(nums) == _brute_left(nums), nums
+
+
+def test_circ_random():
+    rng = random.Random(99)
+    for _ in range(2000):
+        n = rng.randint(0, 40)
+        nums = [rng.randint(-5, 5) for _ in range(n)]
+        assert next_greater_circular(nums) == _brute_circ(nums), nums
+
+
+def test_edge():
+    assert next_greater_right([]) == []
+    assert next_greater_left([]) == []
+    assert next_greater_circular([]) == []
+    assert next_greater_right([5]) == [-1]
+    assert next_greater_circular([1, 2, 3]) == [1, 2, -1]