ndim_ds 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ndim_ds-0.1.0/PKG-INFO +80 -0
- ndim_ds-0.1.0/README.md +73 -0
- ndim_ds-0.1.0/pyproject.toml +25 -0
- ndim_ds-0.1.0/setup.cfg +4 -0
- ndim_ds-0.1.0/src/__init__.py +0 -0
- ndim_ds-0.1.0/src/dynamic_n_dim_diff_fenwick_tree.py +101 -0
- ndim_ds-0.1.0/src/dynamic_n_dim_fenwick_tree.py +111 -0
- ndim_ds-0.1.0/src/dynamic_n_dim_range_fenwick_tree.py +145 -0
- ndim_ds-0.1.0/src/dynamic_n_dim_seg_tree.py +106 -0
- ndim_ds-0.1.0/src/ndim_ds.egg-info/PKG-INFO +80 -0
- ndim_ds-0.1.0/src/ndim_ds.egg-info/SOURCES.txt +19 -0
- ndim_ds-0.1.0/src/ndim_ds.egg-info/dependency_links.txt +1 -0
- ndim_ds-0.1.0/src/ndim_ds.egg-info/top_level.txt +7 -0
- ndim_ds-0.1.0/src/static_n_dim_difference_array.py +101 -0
- ndim_ds-0.1.0/src/static_n_dim_prefix_sum.py +98 -0
- ndim_ds-0.1.0/test/test_dynamic_diff_fenwick.py +51 -0
- ndim_ds-0.1.0/test/test_dynamic_fenwick_tree.py +56 -0
- ndim_ds-0.1.0/test/test_dynamic_range_fenwick.py +55 -0
- ndim_ds-0.1.0/test/test_dynamic_seg_tree.py +61 -0
- ndim_ds-0.1.0/test/test_static_difference_array.py +69 -0
- ndim_ds-0.1.0/test/test_static_prefix_sum.py +68 -0
ndim_ds-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ndim_ds
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A collection of optimized N-dimensional data structures in Python
|
|
5
|
+
Requires-Python: >=3.8
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
|
|
8
|
+
# N-Dimensional Data Structures
|
|
9
|
+
|
|
10
|
+
A blazing-fast, polyglot (Python & C++) library providing generalized $N$-dimensional data structures.
|
|
11
|
+
|
|
12
|
+
Whether you are working in Python for data science and backend engineering, or you need highly-optimized header-only C++ templates to copy-paste into competitive programming platforms like Codeforces and LeetCode, this repository has you covered.
|
|
13
|
+
|
|
14
|
+
## Data Structures Included
|
|
15
|
+
All structures support $O(1)$ or $O(\log^N(V))$ operations scaled effortlessly across any arbitrary number of dimensions.
|
|
16
|
+
|
|
17
|
+
1. **Static N-Dim Prefix Sum** (`static_n_dim_prefix_sum`)
|
|
18
|
+
2. **Static N-Dim Difference Array** (`static_n_dim_difference_array`)
|
|
19
|
+
3. **Dynamic N-Dim Fenwick Tree** (`dynamic_n_dim_fenwick_tree`)
|
|
20
|
+
4. **Dynamic N-Dim Difference Fenwick Tree** (`dynamic_n_dim_diff_fenwick_tree`)
|
|
21
|
+
5. **Dynamic N-Dim Range Fenwick Tree** (`dynamic_n_dim_range_fenwick_tree`)
|
|
22
|
+
6. **Dynamic N-Dim Segment Tree** (`dynamic_n_dim_seg_tree`)
|
|
23
|
+
|
|
24
|
+
---
|
|
25
|
+
|
|
26
|
+
## C++ (Competitive Programming)
|
|
27
|
+
|
|
28
|
+
The `cpp/include/ndim/` directory contains highly optimized, header-only C++17 templates designed specifically for **competitive programming**.
|
|
29
|
+
|
|
30
|
+
They use standard `<bits/stdc++.h>` format with `using namespace std;` to ensure they are instantly copy-pasteable into online judges without causing scope or header errors. Memory is maintained via flat, contiguous 1D vectors mapped algebraically to $N$ dimensions, preventing memory scattering and ensuring maximum CPU cache efficiency.
|
|
31
|
+
|
|
32
|
+
### Example Usage
|
|
33
|
+
```cpp
|
|
34
|
+
#include "dynamic_n_dim_seg_tree.hpp"
|
|
35
|
+
|
|
36
|
+
// Example: 3D Segment Tree using std::min
|
|
37
|
+
auto min_func = [](int64_t a, int64_t b) { return min(a, b); };
|
|
38
|
+
int64_t def = 1e18; // infinity
|
|
39
|
+
|
|
40
|
+
// 4x4x4 grid
|
|
41
|
+
DynamicNDimSegTree<int64_t, decltype(min_func)> tree({4, 4, 4}, min_func, def);
|
|
42
|
+
|
|
43
|
+
tree.update({1, 1, 1}, 5);
|
|
44
|
+
tree.update({2, 2, 2}, 10);
|
|
45
|
+
|
|
46
|
+
int64_t val = tree.query_range({0, 0, 0}, {3, 3, 3}); // 5
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
---
|
|
50
|
+
|
|
51
|
+
## Python (Data Science / General)
|
|
52
|
+
|
|
53
|
+
The `python/src/` directory contains the pure-Python implementations. They use 1D list-flattening mathematics similar to numpy under the hood, but operate purely on standard library primitives, ensuring maximum portability without heavy C-extension dependencies.
|
|
54
|
+
|
|
55
|
+
### Installation
|
|
56
|
+
The library is configured using `uv` via `pyproject.toml`.
|
|
57
|
+
```bash
|
|
58
|
+
cd python
|
|
59
|
+
uv sync
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Example Usage
|
|
63
|
+
```python
|
|
64
|
+
from dynamic_n_dim_range_fenwick_tree import DynamicNDimRangeFenwickTree
|
|
65
|
+
|
|
66
|
+
# 10x10x10 grid
|
|
67
|
+
tree = DynamicNDimRangeFenwickTree([10, 10, 10])
|
|
68
|
+
|
|
69
|
+
# Add 50 to the bounding box from (1, 1, 1) to (5, 5, 5)
|
|
70
|
+
tree.add_range([1, 1, 1], [5, 5, 5], 50)
|
|
71
|
+
|
|
72
|
+
# Query the volume sum from (0, 0, 0) to (3, 3, 3)
|
|
73
|
+
total = tree.query_range([0, 0, 0], [3, 3, 3])
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
## Testing
|
|
77
|
+
|
|
78
|
+
Both ecosystems are rigorously tested with aggressive $10^3$ iteration stress tests across $5$-dimensional constraints.
|
|
79
|
+
- **Python**: Uses `pytest`. Run `uv run pytest test/`
|
|
80
|
+
- **C++**: Uses `doctest`. Compile the files in `cpp/test/` using `g++ -std=c++17` and run the resulting executables.
|
ndim_ds-0.1.0/README.md
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
# N-Dimensional Data Structures
|
|
2
|
+
|
|
3
|
+
A blazing-fast, polyglot (Python & C++) library providing generalized $N$-dimensional data structures.
|
|
4
|
+
|
|
5
|
+
Whether you are working in Python for data science and backend engineering, or you need highly-optimized header-only C++ templates to copy-paste into competitive programming platforms like Codeforces and LeetCode, this repository has you covered.
|
|
6
|
+
|
|
7
|
+
## Data Structures Included
|
|
8
|
+
All structures support $O(1)$ or $O(\log^N(V))$ operations scaled effortlessly across any arbitrary number of dimensions.
|
|
9
|
+
|
|
10
|
+
1. **Static N-Dim Prefix Sum** (`static_n_dim_prefix_sum`)
|
|
11
|
+
2. **Static N-Dim Difference Array** (`static_n_dim_difference_array`)
|
|
12
|
+
3. **Dynamic N-Dim Fenwick Tree** (`dynamic_n_dim_fenwick_tree`)
|
|
13
|
+
4. **Dynamic N-Dim Difference Fenwick Tree** (`dynamic_n_dim_diff_fenwick_tree`)
|
|
14
|
+
5. **Dynamic N-Dim Range Fenwick Tree** (`dynamic_n_dim_range_fenwick_tree`)
|
|
15
|
+
6. **Dynamic N-Dim Segment Tree** (`dynamic_n_dim_seg_tree`)
|
|
16
|
+
|
|
17
|
+
---
|
|
18
|
+
|
|
19
|
+
## C++ (Competitive Programming)
|
|
20
|
+
|
|
21
|
+
The `cpp/include/ndim/` directory contains highly optimized, header-only C++17 templates designed specifically for **competitive programming**.
|
|
22
|
+
|
|
23
|
+
They use standard `<bits/stdc++.h>` format with `using namespace std;` to ensure they are instantly copy-pasteable into online judges without causing scope or header errors. Memory is maintained via flat, contiguous 1D vectors mapped algebraically to $N$ dimensions, preventing memory scattering and ensuring maximum CPU cache efficiency.
|
|
24
|
+
|
|
25
|
+
### Example Usage
|
|
26
|
+
```cpp
|
|
27
|
+
#include "dynamic_n_dim_seg_tree.hpp"
|
|
28
|
+
|
|
29
|
+
// Example: 3D Segment Tree using std::min
|
|
30
|
+
auto min_func = [](int64_t a, int64_t b) { return min(a, b); };
|
|
31
|
+
int64_t def = 1e18; // infinity
|
|
32
|
+
|
|
33
|
+
// 4x4x4 grid
|
|
34
|
+
DynamicNDimSegTree<int64_t, decltype(min_func)> tree({4, 4, 4}, min_func, def);
|
|
35
|
+
|
|
36
|
+
tree.update({1, 1, 1}, 5);
|
|
37
|
+
tree.update({2, 2, 2}, 10);
|
|
38
|
+
|
|
39
|
+
int64_t val = tree.query_range({0, 0, 0}, {3, 3, 3}); // 5
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
---
|
|
43
|
+
|
|
44
|
+
## Python (Data Science / General)
|
|
45
|
+
|
|
46
|
+
The `python/src/` directory contains the pure-Python implementations. They use 1D list-flattening mathematics similar to numpy under the hood, but operate purely on standard library primitives, ensuring maximum portability without heavy C-extension dependencies.
|
|
47
|
+
|
|
48
|
+
### Installation
|
|
49
|
+
The library is configured using `uv` via `pyproject.toml`.
|
|
50
|
+
```bash
|
|
51
|
+
cd python
|
|
52
|
+
uv sync
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
### Example Usage
|
|
56
|
+
```python
|
|
57
|
+
from dynamic_n_dim_range_fenwick_tree import DynamicNDimRangeFenwickTree
|
|
58
|
+
|
|
59
|
+
# 10x10x10 grid
|
|
60
|
+
tree = DynamicNDimRangeFenwickTree([10, 10, 10])
|
|
61
|
+
|
|
62
|
+
# Add 50 to the bounding box from (1, 1, 1) to (5, 5, 5)
|
|
63
|
+
tree.add_range([1, 1, 1], [5, 5, 5], 50)
|
|
64
|
+
|
|
65
|
+
# Query the volume sum from (0, 0, 0) to (3, 3, 3)
|
|
66
|
+
total = tree.query_range([0, 0, 0], [3, 3, 3])
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
## Testing
|
|
70
|
+
|
|
71
|
+
Both ecosystems are rigorously tested with aggressive $10^3$ iteration stress tests across $5$-dimensional constraints.
|
|
72
|
+
- **Python**: Uses `pytest`. Run `uv run pytest test/`
|
|
73
|
+
- **C++**: Uses `doctest`. Compile the files in `cpp/test/` using `g++ -std=c++17` and run the resulting executables.
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "ndim_ds"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "A collection of optimized N-dimensional data structures in Python"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.8"
|
|
7
|
+
dependencies = []
|
|
8
|
+
|
|
9
|
+
[tool.uv]
|
|
10
|
+
dev-dependencies = [
|
|
11
|
+
"pytest>=8.0.0",
|
|
12
|
+
"ruff>=0.4.0",
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
[tool.pytest.ini_options]
|
|
16
|
+
pythonpath = ["src"]
|
|
17
|
+
testpaths = ["test"]
|
|
18
|
+
|
|
19
|
+
[tool.ruff]
|
|
20
|
+
line-length = 100
|
|
21
|
+
target-version = "py38"
|
|
22
|
+
|
|
23
|
+
[tool.ruff.lint]
|
|
24
|
+
select = ["E", "F", "W", "I"]
|
|
25
|
+
ignore = []
|
ndim_ds-0.1.0/setup.cfg
ADDED
|
File without changes
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class DynamicNDimDiffFenwickTree:
|
|
5
|
+
def __init__(self, dims):
|
|
6
|
+
"""
|
|
7
|
+
Initializes an N-dimensional dynamic difference array powered by a Fenwick Tree.
|
|
8
|
+
This structure is optimized for applying values across continuous ranges
|
|
9
|
+
and querying the exact value of any single point dynamically.
|
|
10
|
+
|
|
11
|
+
The grid uses 0-based indexing and allocates memory strictly for the requested
|
|
12
|
+
dimensions to prevent exponential memory bloat.
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
dims (iterable of int): The strict maximum size of the grid in each dimension.
|
|
16
|
+
For example, dims=[10, 10, 10] creates a 3D grid
|
|
17
|
+
allowing coordinates from (0,0,0) to (9,9,9).
|
|
18
|
+
"""
|
|
19
|
+
self.n = len(dims)
|
|
20
|
+
self.dims = tuple(dims)
|
|
21
|
+
|
|
22
|
+
self.strides = [1] * self.n
|
|
23
|
+
for i in range(self.n - 2, -1, -1):
|
|
24
|
+
self.strides[i] = self.strides[i + 1] * self.dims[i + 1]
|
|
25
|
+
|
|
26
|
+
self.total_size = self.strides[0] * self.dims[0]
|
|
27
|
+
self.arr = [0] * self.total_size
|
|
28
|
+
|
|
29
|
+
def _add_point(self, coords, val):
|
|
30
|
+
"""
|
|
31
|
+
Internal helper method to apply a value to a specific point in the Fenwick Tree.
|
|
32
|
+
Uses 1-based bitwise traversal mapped safely to the 0-based flat array.
|
|
33
|
+
"""
|
|
34
|
+
dim_indices = []
|
|
35
|
+
for d in range(self.n):
|
|
36
|
+
idx_list = []
|
|
37
|
+
i = coords[d] + 1
|
|
38
|
+
while i <= self.dims[d]:
|
|
39
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
40
|
+
i += i & (-i)
|
|
41
|
+
dim_indices.append(idx_list)
|
|
42
|
+
|
|
43
|
+
for offsets in itertools.product(*dim_indices):
|
|
44
|
+
self.arr[sum(offsets)] += val
|
|
45
|
+
|
|
46
|
+
def add_range(self, x_coords, y_coords, val):
|
|
47
|
+
"""
|
|
48
|
+
Dynamically adds a value to all elements within an N-dimensional bounding box.
|
|
49
|
+
Places 2^N markers at the boundaries using the Inclusion-Exclusion principle.
|
|
50
|
+
|
|
51
|
+
Args:
|
|
52
|
+
x_coords (iterable of int): 0-based starting indices (lower bounds) of the bounding box.
|
|
53
|
+
y_coords (iterable of int): 0-based ending indices (inclusive upper bounds).
|
|
54
|
+
val (int or float): The value to add to every point inside the specified range.
|
|
55
|
+
"""
|
|
56
|
+
for mask in range(1 << self.n):
|
|
57
|
+
p_coords = [0] * self.n
|
|
58
|
+
sign = 1
|
|
59
|
+
valid = True
|
|
60
|
+
|
|
61
|
+
for d in range(self.n):
|
|
62
|
+
if (mask >> d) & 1:
|
|
63
|
+
c = y_coords[d] + 1
|
|
64
|
+
if c >= self.dims[d]:
|
|
65
|
+
valid = False
|
|
66
|
+
break
|
|
67
|
+
p_coords[d] = c
|
|
68
|
+
sign = -sign
|
|
69
|
+
else:
|
|
70
|
+
p_coords[d] = x_coords[d]
|
|
71
|
+
|
|
72
|
+
if valid:
|
|
73
|
+
self._add_point(p_coords, sign * val)
|
|
74
|
+
|
|
75
|
+
def query_point(self, coords):
|
|
76
|
+
"""
|
|
77
|
+
Retrieves the exact, current value at a specific point in the N-dimensional grid.
|
|
78
|
+
Since this tree stores a difference array, a point's value is calculated
|
|
79
|
+
by summing all the difference markers from the origin up to the given coordinates.
|
|
80
|
+
|
|
81
|
+
Args:
|
|
82
|
+
coords (iterable of int): 0-based coordinates of the exact point to query.
|
|
83
|
+
|
|
84
|
+
Returns:
|
|
85
|
+
int or float: The current cumulative value at the specified point.
|
|
86
|
+
"""
|
|
87
|
+
dim_indices = []
|
|
88
|
+
for d in range(self.n):
|
|
89
|
+
idx_list = []
|
|
90
|
+
i = coords[d] + 1
|
|
91
|
+
i = min(i, self.dims[d])
|
|
92
|
+
while i > 0:
|
|
93
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
94
|
+
i -= i & (-i)
|
|
95
|
+
dim_indices.append(idx_list)
|
|
96
|
+
|
|
97
|
+
res = 0
|
|
98
|
+
for offsets in itertools.product(*dim_indices):
|
|
99
|
+
res += self.arr[sum(offsets)]
|
|
100
|
+
|
|
101
|
+
return res
|
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class DynamicNDimFenwickTree:
|
|
5
|
+
def __init__(self, dims):
|
|
6
|
+
"""
|
|
7
|
+
Initializes an N-dimensional dynamic grid (Fenwick Tree / Binary Indexed Tree)
|
|
8
|
+
capable of processing point updates and range queries in logarithmic time.
|
|
9
|
+
|
|
10
|
+
The grid uses 0-based indexing and allocates memory strictly for the requested
|
|
11
|
+
dimensions, ensuring no wasted padding space.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
dims (iterable of int): The strict maximum size of the grid in each dimension.
|
|
15
|
+
For example, dims=[5, 5] allows coordinates from (0,0) to (4,4).
|
|
16
|
+
"""
|
|
17
|
+
self.n = len(dims)
|
|
18
|
+
self.dims = tuple(dims)
|
|
19
|
+
|
|
20
|
+
self.strides = [1] * self.n
|
|
21
|
+
for i in range(self.n - 2, -1, -1):
|
|
22
|
+
self.strides[i] = self.strides[i + 1] * self.dims[i + 1]
|
|
23
|
+
|
|
24
|
+
self.total_size = self.strides[0] * self.dims[0]
|
|
25
|
+
self.arr = [0] * self.total_size
|
|
26
|
+
|
|
27
|
+
def add(self, coords, val):
|
|
28
|
+
"""
|
|
29
|
+
Adds a value to a specific single point in the N-dimensional grid.
|
|
30
|
+
This operation updates the internal tree structure dynamically.
|
|
31
|
+
|
|
32
|
+
Args:
|
|
33
|
+
coords (iterable of int): 0-based coordinate indices of the point to update.
|
|
34
|
+
Must be of length N.
|
|
35
|
+
val (int or float): The value to add to the specified point.
|
|
36
|
+
"""
|
|
37
|
+
dim_indices = []
|
|
38
|
+
for d in range(self.n):
|
|
39
|
+
idx_list = []
|
|
40
|
+
i = coords[d] + 1
|
|
41
|
+
while i <= self.dims[d]:
|
|
42
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
43
|
+
i += i & (-i)
|
|
44
|
+
dim_indices.append(idx_list)
|
|
45
|
+
|
|
46
|
+
for offsets in itertools.product(*dim_indices):
|
|
47
|
+
self.arr[sum(offsets)] += val
|
|
48
|
+
|
|
49
|
+
def query_prefix(self, coords):
|
|
50
|
+
"""
|
|
51
|
+
Calculates the dynamic sum of all elements from the origin (0, 0, ...)
|
|
52
|
+
up to and including the given coordinates.
|
|
53
|
+
|
|
54
|
+
Args:
|
|
55
|
+
coords (iterable of int): 0-based ending indices (inclusive).
|
|
56
|
+
Must be of length N.
|
|
57
|
+
|
|
58
|
+
Returns:
|
|
59
|
+
int or float: The prefix sum from the origin to the specified coordinates.
|
|
60
|
+
"""
|
|
61
|
+
dim_indices = []
|
|
62
|
+
for d in range(self.n):
|
|
63
|
+
idx_list = []
|
|
64
|
+
i = coords[d] + 1
|
|
65
|
+
i = min(i, self.dims[d])
|
|
66
|
+
while i > 0:
|
|
67
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
68
|
+
i -= i & (-i)
|
|
69
|
+
dim_indices.append(idx_list)
|
|
70
|
+
|
|
71
|
+
res = 0
|
|
72
|
+
for offsets in itertools.product(*dim_indices):
|
|
73
|
+
res += self.arr[sum(offsets)]
|
|
74
|
+
|
|
75
|
+
return res
|
|
76
|
+
|
|
77
|
+
def query_range(self, x_coords, y_coords):
|
|
78
|
+
"""
|
|
79
|
+
Calculates the dynamic sum of all elements within an N-dimensional bounding box.
|
|
80
|
+
Uses the Inclusion-Exclusion principle over the $2^N$ vertices of the hyper-rectangle.
|
|
81
|
+
|
|
82
|
+
Args:
|
|
83
|
+
x_coords (iterable of int): 0-based starting indices (lower bounds)
|
|
84
|
+
of the bounding box.
|
|
85
|
+
y_coords (iterable of int): 0-based ending indices (inclusive upper bounds)
|
|
86
|
+
of the bounding box.
|
|
87
|
+
|
|
88
|
+
Returns:
|
|
89
|
+
int or float: The sum of the elements within the specified bounding box.
|
|
90
|
+
"""
|
|
91
|
+
ans = 0
|
|
92
|
+
for mask in range(1 << self.n):
|
|
93
|
+
q_coords = [0] * self.n
|
|
94
|
+
sign = 1
|
|
95
|
+
valid = True
|
|
96
|
+
|
|
97
|
+
for d in range(self.n):
|
|
98
|
+
if (mask >> d) & 1:
|
|
99
|
+
c = x_coords[d] - 1
|
|
100
|
+
if c < 0:
|
|
101
|
+
valid = False
|
|
102
|
+
break
|
|
103
|
+
q_coords[d] = c
|
|
104
|
+
sign = -sign
|
|
105
|
+
else:
|
|
106
|
+
q_coords[d] = y_coords[d]
|
|
107
|
+
|
|
108
|
+
if valid:
|
|
109
|
+
ans += sign * self.query_prefix(q_coords)
|
|
110
|
+
|
|
111
|
+
return ans
|
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class DynamicNDimRangeFenwickTree:
|
|
5
|
+
def __init__(self, dims):
|
|
6
|
+
"""
|
|
7
|
+
Initializes an N-dimensional Fenwick Tree capable of processing BOTH
|
|
8
|
+
dynamic range updates and dynamic range queries in logarithmic time.
|
|
9
|
+
|
|
10
|
+
This relies on the "2^N Algebraic Expansion Trick", maintaining 2^N
|
|
11
|
+
virtual trees packed tightly into a single array to compute spatial volumes.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
dims (iterable of int): The strict maximum size of the grid in each dimension.
|
|
15
|
+
Memory footprint is O(2^N * D_0 * D_1 * ... * D_n).
|
|
16
|
+
"""
|
|
17
|
+
self.n = len(dims)
|
|
18
|
+
self.dims = tuple(dims)
|
|
19
|
+
self.num_masks = 1 << self.n
|
|
20
|
+
|
|
21
|
+
self.strides = [1] * self.n
|
|
22
|
+
for i in range(self.n - 2, -1, -1):
|
|
23
|
+
self.strides[i] = self.strides[i + 1] * self.dims[i + 1]
|
|
24
|
+
|
|
25
|
+
self.total_size = self.strides[0] * self.dims[0]
|
|
26
|
+
self.arr = [[0] * self.num_masks for _ in range(self.total_size)]
|
|
27
|
+
|
|
28
|
+
def _add_point(self, coords, val):
|
|
29
|
+
"""
|
|
30
|
+
Internal method. Updates all 2^N algebraic states at a specific point.
|
|
31
|
+
"""
|
|
32
|
+
mask_vals = [0] * self.num_masks
|
|
33
|
+
for mask in range(self.num_masks):
|
|
34
|
+
v = val
|
|
35
|
+
for d in range(self.n):
|
|
36
|
+
if (mask >> d) & 1:
|
|
37
|
+
v *= -(coords[d] + 1)
|
|
38
|
+
mask_vals[mask] = v
|
|
39
|
+
|
|
40
|
+
dim_indices = []
|
|
41
|
+
for d in range(self.n):
|
|
42
|
+
idx_list = []
|
|
43
|
+
i = coords[d] + 1
|
|
44
|
+
while i <= self.dims[d]:
|
|
45
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
46
|
+
i += i & (-i)
|
|
47
|
+
dim_indices.append(idx_list)
|
|
48
|
+
|
|
49
|
+
for offsets in itertools.product(*dim_indices):
|
|
50
|
+
idx = sum(offsets)
|
|
51
|
+
for mask in range(self.num_masks):
|
|
52
|
+
self.arr[idx][mask] += mask_vals[mask]
|
|
53
|
+
|
|
54
|
+
def add_range(self, x_coords, y_coords, val):
|
|
55
|
+
"""
|
|
56
|
+
Dynamically adds a value to all elements within an N-dimensional bounding box.
|
|
57
|
+
|
|
58
|
+
Args:
|
|
59
|
+
x_coords (iterable of int): 0-based starting indices (lower bounds).
|
|
60
|
+
y_coords (iterable of int): 0-based ending indices (inclusive upper bounds).
|
|
61
|
+
val (int or float): The value to add to the specified volume.
|
|
62
|
+
"""
|
|
63
|
+
for mask in range(self.num_masks):
|
|
64
|
+
p_coords = [0] * self.n
|
|
65
|
+
sign = 1
|
|
66
|
+
valid = True
|
|
67
|
+
|
|
68
|
+
for d in range(self.n):
|
|
69
|
+
if (mask >> d) & 1:
|
|
70
|
+
c = y_coords[d] + 1
|
|
71
|
+
if c >= self.dims[d]:
|
|
72
|
+
valid = False
|
|
73
|
+
break
|
|
74
|
+
p_coords[d] = c
|
|
75
|
+
sign = -sign
|
|
76
|
+
else:
|
|
77
|
+
p_coords[d] = x_coords[d]
|
|
78
|
+
|
|
79
|
+
if valid:
|
|
80
|
+
self._add_point(p_coords, sign * val)
|
|
81
|
+
|
|
82
|
+
def query_prefix(self, coords):
|
|
83
|
+
"""
|
|
84
|
+
Calculates the spatial volume from the origin up to the given coordinates
|
|
85
|
+
using the 2^N algebraic expansion.
|
|
86
|
+
"""
|
|
87
|
+
clamped_coords = [min(coords[d], self.dims[d] - 1) for d in range(self.n)]
|
|
88
|
+
|
|
89
|
+
dim_indices = []
|
|
90
|
+
for d in range(self.n):
|
|
91
|
+
idx_list = []
|
|
92
|
+
i = clamped_coords[d] + 1
|
|
93
|
+
while i > 0:
|
|
94
|
+
idx_list.append((i - 1) * self.strides[d])
|
|
95
|
+
i -= i & (-i)
|
|
96
|
+
dim_indices.append(idx_list)
|
|
97
|
+
|
|
98
|
+
tree_sums = [0] * self.num_masks
|
|
99
|
+
for offsets in itertools.product(*dim_indices):
|
|
100
|
+
idx = sum(offsets)
|
|
101
|
+
for mask in range(self.num_masks):
|
|
102
|
+
tree_sums[mask] += self.arr[idx][mask]
|
|
103
|
+
|
|
104
|
+
res = 0
|
|
105
|
+
for mask in range(self.num_masks):
|
|
106
|
+
multiplier = 1
|
|
107
|
+
for d in range(self.n):
|
|
108
|
+
if not ((mask >> d) & 1):
|
|
109
|
+
multiplier *= clamped_coords[d] + 2
|
|
110
|
+
res += multiplier * tree_sums[mask]
|
|
111
|
+
|
|
112
|
+
return res
|
|
113
|
+
|
|
114
|
+
def query_range(self, x_coords, y_coords):
|
|
115
|
+
"""
|
|
116
|
+
Calculates the dynamic sum of all elements within an N-dimensional bounding box.
|
|
117
|
+
|
|
118
|
+
Args:
|
|
119
|
+
x_coords (iterable of int): 0-based starting indices.
|
|
120
|
+
y_coords (iterable of int): 0-based ending indices (inclusive).
|
|
121
|
+
|
|
122
|
+
Returns:
|
|
123
|
+
int or float: The dynamic sum within the specified volume.
|
|
124
|
+
"""
|
|
125
|
+
ans = 0
|
|
126
|
+
for mask in range(self.num_masks):
|
|
127
|
+
q_coords = [0] * self.n
|
|
128
|
+
sign = 1
|
|
129
|
+
valid = True
|
|
130
|
+
|
|
131
|
+
for d in range(self.n):
|
|
132
|
+
if (mask >> d) & 1:
|
|
133
|
+
c = x_coords[d] - 1
|
|
134
|
+
if c < 0:
|
|
135
|
+
valid = False
|
|
136
|
+
break
|
|
137
|
+
q_coords[d] = c
|
|
138
|
+
sign = -sign
|
|
139
|
+
else:
|
|
140
|
+
q_coords[d] = y_coords[d]
|
|
141
|
+
|
|
142
|
+
if valid:
|
|
143
|
+
ans += sign * self.query_prefix(q_coords)
|
|
144
|
+
|
|
145
|
+
return ans
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class DynamicNDimSegTree:
|
|
5
|
+
def __init__(self, dims, func=min, default=float("inf")):
|
|
6
|
+
"""
|
|
7
|
+
Initializes an N-dimensional Iterative Segment Tree.
|
|
8
|
+
Perfect for dynamic idempotent operations (Min, Max, GCD).
|
|
9
|
+
|
|
10
|
+
The grid uses 0-based indexing. It leverages an iterative array layout
|
|
11
|
+
to strictly bound memory to exactly 2^N * D_1 * D_2 ... * D_n.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
dims (iterable of int): The strict maximum size of the grid in each dimension.
|
|
15
|
+
func (callable): The function to merge two nodes (e.g., min, max, math.gcd).
|
|
16
|
+
default (int or float): The identity value for the function.
|
|
17
|
+
"""
|
|
18
|
+
self.n = len(dims)
|
|
19
|
+
self.dims = tuple(dims)
|
|
20
|
+
self.func = func
|
|
21
|
+
self.default = default
|
|
22
|
+
|
|
23
|
+
self.sizes = tuple(2 * d for d in dims)
|
|
24
|
+
|
|
25
|
+
self.strides = [1] * self.n
|
|
26
|
+
for i in range(self.n - 2, -1, -1):
|
|
27
|
+
self.strides[i] = self.strides[i + 1] * self.sizes[i + 1]
|
|
28
|
+
|
|
29
|
+
self.total_size = self.strides[0] * self.sizes[0]
|
|
30
|
+
self.arr = [self.default] * self.total_size
|
|
31
|
+
|
|
32
|
+
def update(self, coords, val):
|
|
33
|
+
"""
|
|
34
|
+
Updates a specific point and iteratively recomputes the affected tree hierarchy.
|
|
35
|
+
|
|
36
|
+
Args:
|
|
37
|
+
coords (iterable of int): 0-based coordinate indices of the point.
|
|
38
|
+
val (int or float): The new value for the specified point.
|
|
39
|
+
"""
|
|
40
|
+
leaf_coords = [coords[d] + self.dims[d] for d in range(self.n)]
|
|
41
|
+
|
|
42
|
+
paths = []
|
|
43
|
+
for d in range(self.n):
|
|
44
|
+
path = []
|
|
45
|
+
p = leaf_coords[d]
|
|
46
|
+
while p >= 1:
|
|
47
|
+
path.append(p)
|
|
48
|
+
p >>= 1
|
|
49
|
+
paths.append(path)
|
|
50
|
+
|
|
51
|
+
for p_tuple in itertools.product(*paths):
|
|
52
|
+
idx = 0
|
|
53
|
+
for d in range(self.n):
|
|
54
|
+
idx += p_tuple[d] * self.strides[d]
|
|
55
|
+
|
|
56
|
+
split_dim = -1
|
|
57
|
+
for d in range(self.n):
|
|
58
|
+
if p_tuple[d] != leaf_coords[d]:
|
|
59
|
+
split_dim = d
|
|
60
|
+
break
|
|
61
|
+
|
|
62
|
+
if split_dim == -1:
|
|
63
|
+
self.arr[idx] = val
|
|
64
|
+
else:
|
|
65
|
+
left_idx = idx + p_tuple[split_dim] * self.strides[split_dim]
|
|
66
|
+
right_idx = left_idx + self.strides[split_dim]
|
|
67
|
+
self.arr[idx] = self.func(self.arr[left_idx], self.arr[right_idx])
|
|
68
|
+
|
|
69
|
+
def query_range(self, x_coords, y_coords):
|
|
70
|
+
"""
|
|
71
|
+
Iteratively calculates the aggregate function over an N-dimensional bounding box.
|
|
72
|
+
|
|
73
|
+
Args:
|
|
74
|
+
x_coords (iterable of int): 0-based starting indices (lower bounds).
|
|
75
|
+
y_coords (iterable of int): 0-based ending indices (inclusive upper bounds).
|
|
76
|
+
|
|
77
|
+
Returns:
|
|
78
|
+
int or float: The aggregated result (e.g., minimum value) in the volume.
|
|
79
|
+
"""
|
|
80
|
+
dim_nodes = []
|
|
81
|
+
|
|
82
|
+
for d in range(self.n):
|
|
83
|
+
if x_coords[d] > y_coords[d]:
|
|
84
|
+
return self.default
|
|
85
|
+
|
|
86
|
+
left_idx = x_coords[d] + self.dims[d]
|
|
87
|
+
r = y_coords[d] + self.dims[d] + 1
|
|
88
|
+
nodes = []
|
|
89
|
+
|
|
90
|
+
while left_idx < r:
|
|
91
|
+
if left_idx % 2 == 1:
|
|
92
|
+
nodes.append(left_idx * self.strides[d])
|
|
93
|
+
left_idx += 1
|
|
94
|
+
if r % 2 == 1:
|
|
95
|
+
r -= 1
|
|
96
|
+
nodes.append(r * self.strides[d])
|
|
97
|
+
left_idx >>= 1
|
|
98
|
+
r >>= 1
|
|
99
|
+
|
|
100
|
+
dim_nodes.append(nodes)
|
|
101
|
+
|
|
102
|
+
res = self.default
|
|
103
|
+
for offsets in itertools.product(*dim_nodes):
|
|
104
|
+
res = self.func(res, self.arr[sum(offsets)])
|
|
105
|
+
|
|
106
|
+
return res
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ndim_ds
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A collection of optimized N-dimensional data structures in Python
|
|
5
|
+
Requires-Python: >=3.8
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
|
|
8
|
+
# N-Dimensional Data Structures
|
|
9
|
+
|
|
10
|
+
A blazing-fast, polyglot (Python & C++) library providing generalized $N$-dimensional data structures.
|
|
11
|
+
|
|
12
|
+
Whether you are working in Python for data science and backend engineering, or you need highly-optimized header-only C++ templates to copy-paste into competitive programming platforms like Codeforces and LeetCode, this repository has you covered.
|
|
13
|
+
|
|
14
|
+
## Data Structures Included
|
|
15
|
+
All structures support $O(1)$ or $O(\log^N(V))$ operations scaled effortlessly across any arbitrary number of dimensions.
|
|
16
|
+
|
|
17
|
+
1. **Static N-Dim Prefix Sum** (`static_n_dim_prefix_sum`)
|
|
18
|
+
2. **Static N-Dim Difference Array** (`static_n_dim_difference_array`)
|
|
19
|
+
3. **Dynamic N-Dim Fenwick Tree** (`dynamic_n_dim_fenwick_tree`)
|
|
20
|
+
4. **Dynamic N-Dim Difference Fenwick Tree** (`dynamic_n_dim_diff_fenwick_tree`)
|
|
21
|
+
5. **Dynamic N-Dim Range Fenwick Tree** (`dynamic_n_dim_range_fenwick_tree`)
|
|
22
|
+
6. **Dynamic N-Dim Segment Tree** (`dynamic_n_dim_seg_tree`)
|
|
23
|
+
|
|
24
|
+
---
|
|
25
|
+
|
|
26
|
+
## C++ (Competitive Programming)
|
|
27
|
+
|
|
28
|
+
The `cpp/include/ndim/` directory contains highly optimized, header-only C++17 templates designed specifically for **competitive programming**.
|
|
29
|
+
|
|
30
|
+
They use standard `<bits/stdc++.h>` format with `using namespace std;` to ensure they are instantly copy-pasteable into online judges without causing scope or header errors. Memory is maintained via flat, contiguous 1D vectors mapped algebraically to $N$ dimensions, preventing memory scattering and ensuring maximum CPU cache efficiency.
|
|
31
|
+
|
|
32
|
+
### Example Usage
|
|
33
|
+
```cpp
|
|
34
|
+
#include "dynamic_n_dim_seg_tree.hpp"
|
|
35
|
+
|
|
36
|
+
// Example: 3D Segment Tree using std::min
|
|
37
|
+
auto min_func = [](int64_t a, int64_t b) { return min(a, b); };
|
|
38
|
+
int64_t def = 1e18; // infinity
|
|
39
|
+
|
|
40
|
+
// 4x4x4 grid
|
|
41
|
+
DynamicNDimSegTree<int64_t, decltype(min_func)> tree({4, 4, 4}, min_func, def);
|
|
42
|
+
|
|
43
|
+
tree.update({1, 1, 1}, 5);
|
|
44
|
+
tree.update({2, 2, 2}, 10);
|
|
45
|
+
|
|
46
|
+
int64_t val = tree.query_range({0, 0, 0}, {3, 3, 3}); // 5
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
---
|
|
50
|
+
|
|
51
|
+
## Python (Data Science / General)
|
|
52
|
+
|
|
53
|
+
The `python/src/` directory contains the pure-Python implementations. They use 1D list-flattening mathematics similar to numpy under the hood, but operate purely on standard library primitives, ensuring maximum portability without heavy C-extension dependencies.
|
|
54
|
+
|
|
55
|
+
### Installation
|
|
56
|
+
The library is configured using `uv` via `pyproject.toml`.
|
|
57
|
+
```bash
|
|
58
|
+
cd python
|
|
59
|
+
uv sync
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Example Usage
|
|
63
|
+
```python
|
|
64
|
+
from dynamic_n_dim_range_fenwick_tree import DynamicNDimRangeFenwickTree
|
|
65
|
+
|
|
66
|
+
# 10x10x10 grid
|
|
67
|
+
tree = DynamicNDimRangeFenwickTree([10, 10, 10])
|
|
68
|
+
|
|
69
|
+
# Add 50 to the bounding box from (1, 1, 1) to (5, 5, 5)
|
|
70
|
+
tree.add_range([1, 1, 1], [5, 5, 5], 50)
|
|
71
|
+
|
|
72
|
+
# Query the volume sum from (0, 0, 0) to (3, 3, 3)
|
|
73
|
+
total = tree.query_range([0, 0, 0], [3, 3, 3])
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
## Testing
|
|
77
|
+
|
|
78
|
+
Both ecosystems are rigorously tested with aggressive $10^3$ iteration stress tests across $5$-dimensional constraints.
|
|
79
|
+
- **Python**: Uses `pytest`. Run `uv run pytest test/`
|
|
80
|
+
- **C++**: Uses `doctest`. Compile the files in `cpp/test/` using `g++ -std=c++17` and run the resulting executables.
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
src/__init__.py
|
|
4
|
+
src/dynamic_n_dim_diff_fenwick_tree.py
|
|
5
|
+
src/dynamic_n_dim_fenwick_tree.py
|
|
6
|
+
src/dynamic_n_dim_range_fenwick_tree.py
|
|
7
|
+
src/dynamic_n_dim_seg_tree.py
|
|
8
|
+
src/static_n_dim_difference_array.py
|
|
9
|
+
src/static_n_dim_prefix_sum.py
|
|
10
|
+
src/ndim_ds.egg-info/PKG-INFO
|
|
11
|
+
src/ndim_ds.egg-info/SOURCES.txt
|
|
12
|
+
src/ndim_ds.egg-info/dependency_links.txt
|
|
13
|
+
src/ndim_ds.egg-info/top_level.txt
|
|
14
|
+
test/test_dynamic_diff_fenwick.py
|
|
15
|
+
test/test_dynamic_fenwick_tree.py
|
|
16
|
+
test/test_dynamic_range_fenwick.py
|
|
17
|
+
test/test_dynamic_seg_tree.py
|
|
18
|
+
test/test_static_difference_array.py
|
|
19
|
+
test/test_static_prefix_sum.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
class StaticNDimDifferenceArray:
|
|
2
|
+
def __init__(self, dims):
|
|
3
|
+
"""
|
|
4
|
+
Initializes an N-dimensional grid of zeros with the specified dimensions.
|
|
5
|
+
|
|
6
|
+
Args:
|
|
7
|
+
dims (iterable of int): The size of the grid in each dimension.
|
|
8
|
+
"""
|
|
9
|
+
self.n = len(dims)
|
|
10
|
+
self.dims = tuple(dims)
|
|
11
|
+
|
|
12
|
+
self.strides = [1] * self.n
|
|
13
|
+
for i in range(self.n - 2, -1, -1):
|
|
14
|
+
self.strides[i] = self.strides[i + 1] * self.dims[i + 1]
|
|
15
|
+
|
|
16
|
+
self.total_size = self.strides[0] * self.dims[0]
|
|
17
|
+
self.arr = [0] * self.total_size
|
|
18
|
+
self.is_swept = False
|
|
19
|
+
|
|
20
|
+
def add_range(self, x_coords, y_coords, val):
|
|
21
|
+
"""
|
|
22
|
+
Adds a value to all elements within an N-dimensional bounding box.
|
|
23
|
+
This operation is deferred and will not be queryable until sweep() is called.
|
|
24
|
+
|
|
25
|
+
Args:
|
|
26
|
+
x_coords (iterable of int): 0-based starting indices of the bounding box
|
|
27
|
+
for each dimension.
|
|
28
|
+
y_coords (iterable of int): 0-based ending indices (inclusive) of the bounding box.
|
|
29
|
+
val (int or float): The value to add to the specified range.
|
|
30
|
+
"""
|
|
31
|
+
if self.is_swept:
|
|
32
|
+
raise RuntimeError("Cannot add after sweep.")
|
|
33
|
+
|
|
34
|
+
for mask in range(1 << self.n):
|
|
35
|
+
idx = 0
|
|
36
|
+
sign = 1
|
|
37
|
+
valid = True
|
|
38
|
+
for d in range(self.n):
|
|
39
|
+
if (mask >> d) & 1:
|
|
40
|
+
c = y_coords[d] + 1
|
|
41
|
+
if c >= self.dims[d]:
|
|
42
|
+
valid = False
|
|
43
|
+
break
|
|
44
|
+
idx += c * self.strides[d]
|
|
45
|
+
sign = -sign
|
|
46
|
+
else:
|
|
47
|
+
c = x_coords[d]
|
|
48
|
+
idx += c * self.strides[d]
|
|
49
|
+
|
|
50
|
+
if valid:
|
|
51
|
+
self.arr[idx] += sign * val
|
|
52
|
+
|
|
53
|
+
def sweep(self):
|
|
54
|
+
"""
|
|
55
|
+
Processes all accumulated range updates and prepares the grid for range queries.
|
|
56
|
+
This transitions the data structure from an update phase to a query phase.
|
|
57
|
+
Must be called exactly once after all additions and before any queries.
|
|
58
|
+
"""
|
|
59
|
+
for _ in range(2):
|
|
60
|
+
for d in range(self.n):
|
|
61
|
+
stride = self.strides[d]
|
|
62
|
+
dim_size = self.dims[d]
|
|
63
|
+
for i in range(self.total_size):
|
|
64
|
+
if (i // stride) % dim_size > 0:
|
|
65
|
+
self.arr[i] += self.arr[i - stride]
|
|
66
|
+
self.is_swept = True
|
|
67
|
+
|
|
68
|
+
def query_range(self, x_coords, y_coords):
|
|
69
|
+
"""
|
|
70
|
+
Calculates the sum of all elements within an N-dimensional bounding box.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
x_coords (iterable of int): 0-based starting indices of the bounding box.
|
|
74
|
+
y_coords (iterable of int): 0-based ending indices (inclusive) of the bounding box.
|
|
75
|
+
|
|
76
|
+
Returns:
|
|
77
|
+
int or float: The sum of the elements in the specified range.
|
|
78
|
+
"""
|
|
79
|
+
if not self.is_swept:
|
|
80
|
+
raise RuntimeError("Must sweep before query.")
|
|
81
|
+
|
|
82
|
+
ans = 0
|
|
83
|
+
for mask in range(1 << self.n):
|
|
84
|
+
idx = 0
|
|
85
|
+
sign = 1
|
|
86
|
+
valid = True
|
|
87
|
+
for d in range(self.n):
|
|
88
|
+
if (mask >> d) & 1:
|
|
89
|
+
c = x_coords[d] - 1
|
|
90
|
+
if c < 0:
|
|
91
|
+
valid = False
|
|
92
|
+
break
|
|
93
|
+
idx += c * self.strides[d]
|
|
94
|
+
sign = -sign
|
|
95
|
+
else:
|
|
96
|
+
c = y_coords[d]
|
|
97
|
+
idx += c * self.strides[d]
|
|
98
|
+
|
|
99
|
+
if valid:
|
|
100
|
+
ans += sign * self.arr[idx]
|
|
101
|
+
return ans
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
class StaticNDimPrefixSum:
|
|
2
|
+
def __init__(self, dims):
|
|
3
|
+
"""
|
|
4
|
+
Initializes an N-dimensional grid optimized for offline point updates
|
|
5
|
+
followed by rapid range queries using a static prefix sum array.
|
|
6
|
+
|
|
7
|
+
The grid uses 0-based indexing and allocates memory strictly for the requested
|
|
8
|
+
dimensions, ensuring no wasted padding space.
|
|
9
|
+
|
|
10
|
+
Args:
|
|
11
|
+
dims (iterable of int): The strict maximum size of the grid in each dimension.
|
|
12
|
+
For example, dims=[5, 5] allows coordinates from (0,0) to (4,4).
|
|
13
|
+
"""
|
|
14
|
+
self.n = len(dims)
|
|
15
|
+
self.dims = tuple(dims)
|
|
16
|
+
|
|
17
|
+
self.strides = [1] * self.n
|
|
18
|
+
for i in range(self.n - 2, -1, -1):
|
|
19
|
+
self.strides[i] = self.strides[i + 1] * self.dims[i + 1]
|
|
20
|
+
|
|
21
|
+
self.total_size = self.strides[0] * self.dims[0]
|
|
22
|
+
self.arr = [0] * self.total_size
|
|
23
|
+
self.is_swept = False
|
|
24
|
+
|
|
25
|
+
def add(self, coords, val):
|
|
26
|
+
"""
|
|
27
|
+
Adds a value to a specific point in the grid.
|
|
28
|
+
This operation is deferred; the true prefix sum must be built by calling sweep()
|
|
29
|
+
before any range queries can be made.
|
|
30
|
+
|
|
31
|
+
Args:
|
|
32
|
+
coords (iterable of int): 0-based coordinate indices of the point.
|
|
33
|
+
val (int or float): The value to add to the specified point.
|
|
34
|
+
"""
|
|
35
|
+
if self.is_swept:
|
|
36
|
+
raise RuntimeError("Cannot add points after the grid has been swept.")
|
|
37
|
+
|
|
38
|
+
idx = 0
|
|
39
|
+
for d in range(self.n):
|
|
40
|
+
idx += coords[d] * self.strides[d]
|
|
41
|
+
|
|
42
|
+
self.arr[idx] += val
|
|
43
|
+
|
|
44
|
+
def sweep(self):
|
|
45
|
+
"""
|
|
46
|
+
Processes all accumulated point updates and computes the N-dimensional prefix sum.
|
|
47
|
+
This transitions the data structure from an update phase to a query phase.
|
|
48
|
+
Must be called exactly once after all additions and before any queries.
|
|
49
|
+
"""
|
|
50
|
+
for d in range(self.n):
|
|
51
|
+
stride = self.strides[d]
|
|
52
|
+
dim_size = self.dims[d]
|
|
53
|
+
for i in range(self.total_size):
|
|
54
|
+
if (i // stride) % dim_size > 0:
|
|
55
|
+
self.arr[i] += self.arr[i - stride]
|
|
56
|
+
|
|
57
|
+
self.is_swept = True
|
|
58
|
+
|
|
59
|
+
def query_range(self, x_coords, y_coords):
|
|
60
|
+
"""
|
|
61
|
+
Calculates the exact sum of all elements within an N-dimensional bounding box.
|
|
62
|
+
Uses the Inclusion-Exclusion principle over the $2^N$ vertices of the hyper-rectangle.
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
x_coords (iterable of int): 0-based starting indices (lower bounds)
|
|
66
|
+
of the bounding box.
|
|
67
|
+
y_coords (iterable of int): 0-based ending indices (inclusive upper bounds)
|
|
68
|
+
of the bounding box.
|
|
69
|
+
|
|
70
|
+
Returns:
|
|
71
|
+
int or float: The sum of the elements within the specified bounding box.
|
|
72
|
+
"""
|
|
73
|
+
if not self.is_swept:
|
|
74
|
+
raise RuntimeError("Must call sweep() before querying.")
|
|
75
|
+
|
|
76
|
+
ans = 0
|
|
77
|
+
for mask in range(1 << self.n):
|
|
78
|
+
idx = 0
|
|
79
|
+
sign = 1
|
|
80
|
+
valid = True
|
|
81
|
+
|
|
82
|
+
for d in range(self.n):
|
|
83
|
+
if (mask >> d) & 1:
|
|
84
|
+
c = x_coords[d] - 1
|
|
85
|
+
|
|
86
|
+
if c < 0:
|
|
87
|
+
valid = False
|
|
88
|
+
break
|
|
89
|
+
idx += c * self.strides[d]
|
|
90
|
+
sign = -sign
|
|
91
|
+
else:
|
|
92
|
+
c = y_coords[d]
|
|
93
|
+
idx += c * self.strides[d]
|
|
94
|
+
|
|
95
|
+
if valid:
|
|
96
|
+
ans += sign * self.arr[idx]
|
|
97
|
+
|
|
98
|
+
return ans
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import random
|
|
3
|
+
|
|
4
|
+
from dynamic_n_dim_diff_fenwick_tree import DynamicNDimDiffFenwickTree
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def test_dynamic_diff_fenwick_2d():
|
|
8
|
+
tree = DynamicNDimDiffFenwickTree([5, 5])
|
|
9
|
+
tree.add_range([1, 1], [2, 2], 5)
|
|
10
|
+
assert tree.query_point([0, 0]) == 0
|
|
11
|
+
assert tree.query_point([1, 1]) == 5
|
|
12
|
+
assert tree.query_point([2, 2]) == 5
|
|
13
|
+
assert tree.query_point([3, 3]) == 0
|
|
14
|
+
assert tree.query_point([1, 2]) == 5
|
|
15
|
+
tree.add_range([2, 2], [3, 3], 10)
|
|
16
|
+
assert tree.query_point([1, 1]) == 5
|
|
17
|
+
assert tree.query_point([2, 2]) == 15
|
|
18
|
+
assert tree.query_point([3, 3]) == 10
|
|
19
|
+
assert tree.query_point([4, 4]) == 0
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_dynamic_diff_fenwick_3d():
|
|
23
|
+
tree = DynamicNDimDiffFenwickTree([4, 4, 4])
|
|
24
|
+
tree.add_range([0, 0, 0], [2, 2, 2], 100)
|
|
25
|
+
assert tree.query_point([0, 0, 0]) == 100
|
|
26
|
+
assert tree.query_point([2, 2, 2]) == 100
|
|
27
|
+
assert tree.query_point([3, 3, 3]) == 0
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_dynamic_diff_fenwick_stress():
|
|
31
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
32
|
+
dims = [4, 4, 4, 4, 4]
|
|
33
|
+
tree = DynamicNDimDiffFenwickTree(dims)
|
|
34
|
+
grid_bf = {}
|
|
35
|
+
|
|
36
|
+
random.seed(42)
|
|
37
|
+
|
|
38
|
+
for _ in range(1000):
|
|
39
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
40
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
41
|
+
val = random.randint(1, 100)
|
|
42
|
+
tree.add_range(x_coords, y_coords, val)
|
|
43
|
+
|
|
44
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
45
|
+
for c in itertools.product(*ranges):
|
|
46
|
+
grid_bf[c] = grid_bf.get(c, 0) + val
|
|
47
|
+
|
|
48
|
+
for _ in range(1000):
|
|
49
|
+
coords = [random.randint(0, d - 1) for d in dims]
|
|
50
|
+
expected = grid_bf.get(tuple(coords), 0)
|
|
51
|
+
assert tree.query_point(coords) == expected
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import random
|
|
3
|
+
|
|
4
|
+
from dynamic_n_dim_fenwick_tree import DynamicNDimFenwickTree
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def test_dynamic_fenwick_tree_2d():
|
|
8
|
+
tree = DynamicNDimFenwickTree([5, 5])
|
|
9
|
+
tree.add([1, 1], 10)
|
|
10
|
+
tree.add([2, 2], 5)
|
|
11
|
+
tree.add([3, 3], 15)
|
|
12
|
+
assert tree.query_range([1, 1], [1, 1]) == 10
|
|
13
|
+
assert tree.query_range([0, 0], [2, 2]) == 15
|
|
14
|
+
assert tree.query_range([1, 1], [4, 4]) == 30
|
|
15
|
+
assert tree.query_range([2, 2], [3, 3]) == 20
|
|
16
|
+
assert tree.query_range([0, 0], [0, 0]) == 0
|
|
17
|
+
tree.add([2, 2], 5)
|
|
18
|
+
assert tree.query_range([2, 2], [3, 3]) == 25
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def test_dynamic_fenwick_tree_3d():
|
|
22
|
+
tree = DynamicNDimFenwickTree([3, 3, 3])
|
|
23
|
+
tree.add([0, 0, 0], 1)
|
|
24
|
+
tree.add([1, 1, 1], 2)
|
|
25
|
+
tree.add([2, 2, 2], 3)
|
|
26
|
+
assert tree.query_range([0, 0, 0], [2, 2, 2]) == 6
|
|
27
|
+
assert tree.query_range([0, 0, 0], [1, 1, 1]) == 3
|
|
28
|
+
assert tree.query_range([1, 1, 1], [2, 2, 2]) == 5
|
|
29
|
+
assert tree.query_prefix([2, 2, 2]) == 6
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def test_dynamic_fenwick_tree_stress():
|
|
33
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
34
|
+
dims = [4, 4, 4, 4, 4]
|
|
35
|
+
tree = DynamicNDimFenwickTree(dims)
|
|
36
|
+
grid_bf = {}
|
|
37
|
+
|
|
38
|
+
random.seed(42)
|
|
39
|
+
|
|
40
|
+
for _ in range(1000):
|
|
41
|
+
coords = [random.randint(0, d - 1) for d in dims]
|
|
42
|
+
val = random.randint(1, 100)
|
|
43
|
+
tree.add(coords, val)
|
|
44
|
+
c_tuple = tuple(coords)
|
|
45
|
+
grid_bf[c_tuple] = grid_bf.get(c_tuple, 0) + val
|
|
46
|
+
|
|
47
|
+
for _ in range(1000):
|
|
48
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
49
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
50
|
+
|
|
51
|
+
expected = 0
|
|
52
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
53
|
+
for c in itertools.product(*ranges):
|
|
54
|
+
expected += grid_bf.get(c, 0)
|
|
55
|
+
|
|
56
|
+
assert tree.query_range(x_coords, y_coords) == expected
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import random
|
|
3
|
+
|
|
4
|
+
from dynamic_n_dim_range_fenwick_tree import DynamicNDimRangeFenwickTree
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def test_dynamic_range_fenwick_2d():
|
|
8
|
+
tree = DynamicNDimRangeFenwickTree([4, 4])
|
|
9
|
+
tree.add_range([0, 0], [1, 1], 5)
|
|
10
|
+
assert tree.query_range([0, 0], [0, 0]) == 5
|
|
11
|
+
assert tree.query_range([1, 1], [1, 1]) == 5
|
|
12
|
+
assert tree.query_range([0, 0], [1, 1]) == 20
|
|
13
|
+
assert tree.query_range([0, 0], [3, 3]) == 20
|
|
14
|
+
tree.add_range([1, 1], [2, 2], 10)
|
|
15
|
+
assert tree.query_range([1, 1], [1, 1]) == 15
|
|
16
|
+
assert tree.query_range([0, 0], [3, 3]) == 60
|
|
17
|
+
assert tree.query_prefix([1, 1]) == 30
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_dynamic_range_fenwick_3d():
|
|
21
|
+
tree = DynamicNDimRangeFenwickTree([3, 3, 3])
|
|
22
|
+
tree.add_range([0, 0, 0], [1, 1, 1], 1)
|
|
23
|
+
assert tree.query_range([0, 0, 0], [0, 0, 0]) == 1
|
|
24
|
+
assert tree.query_range([0, 0, 0], [1, 1, 1]) == 8
|
|
25
|
+
assert tree.query_range([0, 0, 0], [2, 2, 2]) == 8
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_dynamic_range_fenwick_stress():
|
|
29
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
30
|
+
dims = [4, 4, 4, 4, 4]
|
|
31
|
+
tree = DynamicNDimRangeFenwickTree(dims)
|
|
32
|
+
grid_bf = {}
|
|
33
|
+
|
|
34
|
+
random.seed(42)
|
|
35
|
+
|
|
36
|
+
for _ in range(1000):
|
|
37
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
38
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
39
|
+
val = random.randint(1, 100)
|
|
40
|
+
tree.add_range(x_coords, y_coords, val)
|
|
41
|
+
|
|
42
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
43
|
+
for c in itertools.product(*ranges):
|
|
44
|
+
grid_bf[c] = grid_bf.get(c, 0) + val
|
|
45
|
+
|
|
46
|
+
for _ in range(1000):
|
|
47
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
48
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
49
|
+
|
|
50
|
+
expected = 0
|
|
51
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
52
|
+
for c in itertools.product(*ranges):
|
|
53
|
+
expected += grid_bf.get(c, 0)
|
|
54
|
+
|
|
55
|
+
assert tree.query_range(x_coords, y_coords) == expected
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import math
|
|
3
|
+
import random
|
|
4
|
+
|
|
5
|
+
from dynamic_n_dim_seg_tree import DynamicNDimSegTree
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_dynamic_seg_tree_2d():
|
|
9
|
+
tree = DynamicNDimSegTree([4, 4], func=min, default=float("inf"))
|
|
10
|
+
tree.update([0, 0], 10)
|
|
11
|
+
tree.update([1, 1], 5)
|
|
12
|
+
tree.update([2, 2], 15)
|
|
13
|
+
tree.update([3, 3], 20)
|
|
14
|
+
assert tree.query_range([0, 0], [3, 3]) == 5
|
|
15
|
+
assert tree.query_range([0, 0], [0, 0]) == 10
|
|
16
|
+
assert tree.query_range([2, 2], [3, 3]) == 15
|
|
17
|
+
assert tree.query_range([0, 2], [1, 3]) == float("inf")
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_dynamic_seg_tree_gcd():
|
|
21
|
+
tree = DynamicNDimSegTree([2, 2], func=math.gcd, default=0)
|
|
22
|
+
tree.update([0, 0], 24)
|
|
23
|
+
tree.update([1, 1], 36)
|
|
24
|
+
assert tree.query_range([0, 0], [1, 1]) == 12
|
|
25
|
+
assert tree.query_range([0, 0], [0, 0]) == 24
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_dynamic_seg_tree_max():
|
|
29
|
+
tree = DynamicNDimSegTree([3, 3, 3], func=max, default=float("-inf"))
|
|
30
|
+
tree.update([0, 0, 0], 100)
|
|
31
|
+
tree.update([1, 1, 1], 500)
|
|
32
|
+
tree.update([2, 2, 2], 200)
|
|
33
|
+
assert tree.query_range([0, 0, 0], [2, 2, 2]) == 500
|
|
34
|
+
assert tree.query_range([0, 0, 0], [0, 1, 1]) == 100
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_dynamic_seg_tree_stress():
|
|
38
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
39
|
+
dims = [4, 4, 4, 4, 4]
|
|
40
|
+
tree = DynamicNDimSegTree(dims, func=min, default=float("inf"))
|
|
41
|
+
grid_bf = {}
|
|
42
|
+
|
|
43
|
+
random.seed(42)
|
|
44
|
+
|
|
45
|
+
for _ in range(1000):
|
|
46
|
+
coords = [random.randint(0, d - 1) for d in dims]
|
|
47
|
+
val = random.randint(1, 1000)
|
|
48
|
+
tree.update(coords, val)
|
|
49
|
+
c_tuple = tuple(coords)
|
|
50
|
+
grid_bf[c_tuple] = val
|
|
51
|
+
|
|
52
|
+
for _ in range(1000):
|
|
53
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
54
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
55
|
+
|
|
56
|
+
expected = float("inf")
|
|
57
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
58
|
+
for c in itertools.product(*ranges):
|
|
59
|
+
expected = min(expected, grid_bf.get(c, float("inf")))
|
|
60
|
+
|
|
61
|
+
assert tree.query_range(x_coords, y_coords) == expected
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import random
|
|
3
|
+
|
|
4
|
+
import pytest
|
|
5
|
+
from static_n_dim_difference_array import StaticNDimDifferenceArray
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_static_difference_array_2d():
|
|
9
|
+
grid = StaticNDimDifferenceArray([4, 4])
|
|
10
|
+
grid.add_range([1, 1], [2, 2], 5)
|
|
11
|
+
grid.add_range([0, 0], [1, 2], 3)
|
|
12
|
+
grid.sweep()
|
|
13
|
+
assert grid.query_range([0, 0], [0, 0]) == 3
|
|
14
|
+
assert grid.query_range([1, 1], [1, 1]) == 8
|
|
15
|
+
assert grid.query_range([2, 2], [2, 2]) == 5
|
|
16
|
+
assert grid.query_range([3, 3], [3, 3]) == 0
|
|
17
|
+
assert grid.query_range([1, 1], [2, 2]) == 26
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_static_difference_array_3d():
|
|
21
|
+
grid = StaticNDimDifferenceArray([3, 3, 3])
|
|
22
|
+
grid.add_range([0, 0, 0], [1, 1, 1], 10)
|
|
23
|
+
grid.sweep()
|
|
24
|
+
assert grid.query_range([0, 0, 0], [0, 0, 0]) == 10
|
|
25
|
+
assert grid.query_range([1, 1, 1], [1, 1, 1]) == 10
|
|
26
|
+
assert grid.query_range([2, 2, 2], [2, 2, 2]) == 0
|
|
27
|
+
assert grid.query_range([0, 0, 0], [1, 1, 1]) == 80
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_static_difference_array_errors():
|
|
31
|
+
grid = StaticNDimDifferenceArray([2, 2])
|
|
32
|
+
grid.add_range([0, 0], [1, 1], 1)
|
|
33
|
+
with pytest.raises(RuntimeError):
|
|
34
|
+
grid.query_range([0, 0], [1, 1])
|
|
35
|
+
grid.sweep()
|
|
36
|
+
with pytest.raises(RuntimeError):
|
|
37
|
+
grid.add_range([0, 0], [0, 0], 5)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def test_static_difference_array_stress():
|
|
41
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
42
|
+
dims = [4, 4, 4, 4, 4]
|
|
43
|
+
grid = StaticNDimDifferenceArray(dims)
|
|
44
|
+
grid_bf = {}
|
|
45
|
+
|
|
46
|
+
random.seed(42)
|
|
47
|
+
|
|
48
|
+
for _ in range(1000):
|
|
49
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
50
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
51
|
+
val = random.randint(1, 100)
|
|
52
|
+
grid.add_range(x_coords, y_coords, val)
|
|
53
|
+
|
|
54
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
55
|
+
for c in itertools.product(*ranges):
|
|
56
|
+
grid_bf[c] = grid_bf.get(c, 0) + val
|
|
57
|
+
|
|
58
|
+
grid.sweep()
|
|
59
|
+
|
|
60
|
+
for _ in range(1000):
|
|
61
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
62
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
63
|
+
|
|
64
|
+
expected = 0
|
|
65
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
66
|
+
for c in itertools.product(*ranges):
|
|
67
|
+
expected += grid_bf.get(c, 0)
|
|
68
|
+
|
|
69
|
+
assert grid.query_range(x_coords, y_coords) == expected
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import random
|
|
3
|
+
|
|
4
|
+
import pytest
|
|
5
|
+
from static_n_dim_prefix_sum import StaticNDimPrefixSum
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_static_prefix_sum_2d():
|
|
9
|
+
grid = StaticNDimPrefixSum([3, 3])
|
|
10
|
+
grid.add([0, 0], 1)
|
|
11
|
+
grid.add([1, 1], 2)
|
|
12
|
+
grid.add([2, 2], 3)
|
|
13
|
+
grid.sweep()
|
|
14
|
+
assert grid.query_range([0, 0], [0, 0]) == 1
|
|
15
|
+
assert grid.query_range([1, 1], [1, 1]) == 2
|
|
16
|
+
assert grid.query_range([2, 2], [2, 2]) == 3
|
|
17
|
+
assert grid.query_range([0, 0], [2, 2]) == 6
|
|
18
|
+
assert grid.query_range([0, 0], [1, 1]) == 3
|
|
19
|
+
assert grid.query_range([0, 0], [0, 1]) == 1
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_static_prefix_sum_3d():
|
|
23
|
+
grid = StaticNDimPrefixSum([2, 2, 2])
|
|
24
|
+
grid.add([0, 0, 0], 5)
|
|
25
|
+
grid.add([1, 1, 1], 10)
|
|
26
|
+
grid.sweep()
|
|
27
|
+
assert grid.query_range([0, 0, 0], [1, 1, 1]) == 15
|
|
28
|
+
assert grid.query_range([0, 0, 0], [0, 1, 1]) == 5
|
|
29
|
+
assert grid.query_range([1, 1, 1], [1, 1, 1]) == 10
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def test_static_prefix_sum_errors():
|
|
33
|
+
grid = StaticNDimPrefixSum([2, 2])
|
|
34
|
+
grid.add([0, 0], 1)
|
|
35
|
+
with pytest.raises(RuntimeError):
|
|
36
|
+
grid.query_range([0, 0], [1, 1])
|
|
37
|
+
grid.sweep()
|
|
38
|
+
with pytest.raises(RuntimeError):
|
|
39
|
+
grid.add([1, 1], 5)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def test_static_prefix_sum_stress():
|
|
43
|
+
"""Stress test with 5 dimensions (4x4x4x4x4 = 1024) and 1000 calls."""
|
|
44
|
+
dims = [4, 4, 4, 4, 4]
|
|
45
|
+
grid = StaticNDimPrefixSum(dims)
|
|
46
|
+
grid_bf = {}
|
|
47
|
+
|
|
48
|
+
random.seed(42)
|
|
49
|
+
|
|
50
|
+
for _ in range(1000):
|
|
51
|
+
coords = [random.randint(0, d - 1) for d in dims]
|
|
52
|
+
val = random.randint(1, 100)
|
|
53
|
+
grid.add(coords, val)
|
|
54
|
+
c_tuple = tuple(coords)
|
|
55
|
+
grid_bf[c_tuple] = grid_bf.get(c_tuple, 0) + val
|
|
56
|
+
|
|
57
|
+
grid.sweep()
|
|
58
|
+
|
|
59
|
+
for _ in range(1000):
|
|
60
|
+
x_coords = [random.randint(0, d - 1) for d in dims]
|
|
61
|
+
y_coords = [random.randint(x, d - 1) for x, d in zip(x_coords, dims)]
|
|
62
|
+
|
|
63
|
+
expected = 0
|
|
64
|
+
ranges = [range(x, y + 1) for x, y in zip(x_coords, y_coords)]
|
|
65
|
+
for c in itertools.product(*ranges):
|
|
66
|
+
expected += grid_bf.get(c, 0)
|
|
67
|
+
|
|
68
|
+
assert grid.query_range(x_coords, y_coords) == expected
|