sketch_segment_tree.nx source
↩ module page · 204 lines · 6084 B
1// sketch_segment_tree.nx -- segment tree for range sum / min / max queries.
2//
3// Foundational data structure. Array of N elements, supports:
4// point_update(i, value) O(log N)
5// range_sum(lo, hi) O(log N)
6// range_min(lo, hi) O(log N)
7// range_max(lo, hi) O(log N)
8//
9// All three aggregates are maintained per node so callers don't have to
10// pick a single aggregate at construction. Memory: 3 * 2N i64s
11// (sum/min/max per node) for N-element backing array.
12//
13// USE CASES:
14// - online range-sum queries (stocks-OHLC over time-window)
15// - dynamic-programming acceleration (e.g., longest increasing subseq)
16// - inverse-index aggregation
17// - online statistics over sliding indexed ranges
18//
19// API uses HALF-OPEN intervals: range_*(lo, hi) covers indices [lo, hi).
20//
21// LOSSLESS-LANGUAGE DISCIPLINE: all queries EXACT. Production tier.
22// Initialized to:
23// sum = 0
24// min = +MAX (so unwritten slots don't affect min)
25// max = -MAX (similar)
26// Caller responsibility: point_update before query for meaningful min/max
27// over ranges containing unwritten slots.
28
29import "syscalls.nx"
30import "sketch_types.nx"
31
32const NX_ST_MIN_N: i64 = 2
33const NX_ST_MAX_N: i64 = 1000000
34
35const NX_ST_PLUS_INF: i64 = 0x4000000000000000
36const NX_ST_MINUS_INF: i64 = -0x4000000000000000
37
38struct SegmentTree {
39 n: i64,
40 capacity: i64, // next power of 2 >= n
41 sum: *i64, // 2 * capacity entries; index 1 is root
42 min_arr: *i64,
43 max_arr: *i64,
44}
45
46// === construction =================================================
47
48func nx_st_next_pow2(n: i64) -> i64 {
49 var p: i64 = 1
50 while p < n { p = p << 1 }
51 return p
52}
53
54func nx_st_alloc(n: i64) -> *SegmentTree {
55 if n < NX_ST_MIN_N { return 0 as *SegmentTree }
56 if n > NX_ST_MAX_N { return 0 as *SegmentTree }
57 let raw: *u8 = sys_mmap(56)
58 let t: *SegmentTree = raw as *SegmentTree
59 let cap: i64 = nx_st_next_pow2(n)
60 let cells: i64 = 2 * cap
61 t.sum = sys_mmap(cells * 8) as *i64
62 t.min_arr = sys_mmap(cells * 8) as *i64
63 t.max_arr = sys_mmap(cells * 8) as *i64
64 var i: i64 = 0
65 while i < cells {
66 t.sum[i] = 0
67 t.min_arr[i] = NX_ST_PLUS_INF
68 t.max_arr[i] = NX_ST_MINUS_INF
69 i = i + 1
70 }
71 t.n = n
72 t.capacity = cap
73 return t
74}
75
76// === min/max helpers ==============================================
77
78func nx_st_min2(a: i64, b: i64) -> i64 {
79 if a < b { return a }
80 return b
81}
82
83func nx_st_max2(a: i64, b: i64) -> i64 {
84 if a > b { return a }
85 return b
86}
87
88// === point update ================================================
89//
90// Set value at leaf index `idx` to `value`; propagate up to root.
91
92func nx_st_update(t: *SegmentTree, idx: i64, value: i64) -> i64 {
93 if idx < 0 { return -1 }
94 if idx >= t.n { return -1 }
95 // Leaves stored at positions capacity..2*capacity-1.
96 var pos: i64 = t.capacity + idx
97 t.sum[pos] = value
98 t.min_arr[pos] = value
99 t.max_arr[pos] = value
100 pos = pos >> 1
101 while pos > 0 {
102 let lc: i64 = 2 * pos
103 let rc: i64 = 2 * pos + 1
104 t.sum[pos] = t.sum[lc] + t.sum[rc]
105 t.min_arr[pos] = nx_st_min2(t.min_arr[lc], t.min_arr[rc])
106 t.max_arr[pos] = nx_st_max2(t.max_arr[lc], t.max_arr[rc])
107 pos = pos >> 1
108 }
109 return 0
110}
111
112// === range query helpers ==========================================
113//
114// Iterative range queries over [lo, hi) in leaf indices.
115// Translates to leaf positions [capacity+lo, capacity+hi).
116
117func nx_st_range_sum(t: *SegmentTree, lo: i64, hi: i64) -> i64 {
118 if lo < 0 { return 0 }
119 if hi > t.n { return 0 }
120 if lo >= hi { return 0 }
121 var l: i64 = t.capacity + lo
122 var r: i64 = t.capacity + hi
123 var acc: i64 = 0
124 while l < r {
125 if (l & 1) == 1 {
126 acc = acc + t.sum[l]
127 l = l + 1
128 }
129 if (r & 1) == 1 {
130 r = r - 1
131 acc = acc + t.sum[r]
132 }
133 l = l >> 1
134 r = r >> 1
135 }
136 return acc
137}
138
139func nx_st_range_min(t: *SegmentTree, lo: i64, hi: i64) -> i64 {
140 if lo < 0 { return NX_ST_PLUS_INF }
141 if hi > t.n { return NX_ST_PLUS_INF }
142 if lo >= hi { return NX_ST_PLUS_INF }
143 var l: i64 = t.capacity + lo
144 var r: i64 = t.capacity + hi
145 var acc: i64 = NX_ST_PLUS_INF
146 while l < r {
147 if (l & 1) == 1 {
148 acc = nx_st_min2(acc, t.min_arr[l])
149 l = l + 1
150 }
151 if (r & 1) == 1 {
152 r = r - 1
153 acc = nx_st_min2(acc, t.min_arr[r])
154 }
155 l = l >> 1
156 r = r >> 1
157 }
158 return acc
159}
160
161func nx_st_range_max(t: *SegmentTree, lo: i64, hi: i64) -> i64 {
162 if lo < 0 { return NX_ST_MINUS_INF }
163 if hi > t.n { return NX_ST_MINUS_INF }
164 if lo >= hi { return NX_ST_MINUS_INF }
165 var l: i64 = t.capacity + lo
166 var r: i64 = t.capacity + hi
167 var acc: i64 = NX_ST_MINUS_INF
168 while l < r {
169 if (l & 1) == 1 {
170 acc = nx_st_max2(acc, t.max_arr[l])
171 l = l + 1
172 }
173 if (r & 1) == 1 {
174 r = r - 1
175 acc = nx_st_max2(acc, t.max_arr[r])
176 }
177 l = l >> 1
178 r = r >> 1
179 }
180 return acc
181}
182
183// === point read ==================================================
184
185func nx_st_get(t: *SegmentTree, idx: i64) -> i64 {
186 if idx < 0 { return 0 }
187 if idx >= t.n { return 0 }
188 return t.sum[t.capacity + idx]
189}
190
191// === typed query =================================================
192
193func nx_st_query_sum(t: *SegmentTree, lo: i64, hi: i64) -> *ApproxI64 {
194 let s: i64 = nx_st_range_sum(t, lo, hi)
195 return nx_approx_new(s, NX_ENV_ABS, 0, 1000000000,
196 NX_MATURITY_PRODUCTION,
197 NX_ADV_HONEST)
198}
199
200func nx_st_memory_bytes(t: *SegmentTree) -> i64 {
201 return 56 + 3 * 2 * t.capacity * 8
202}
203
204func nx_st_n(t: *SegmentTree) -> i64 { return t.n }