Skip to content

Commit 3e8f30c

Browse files
itsamirhnitsamirhn
authored
Add Lazy Segment Tree (TheAlgorithms#3209)
1 parent d63813e commit 3e8f30c

File tree

2 files changed

+202
-0
lines changed

2 files changed

+202
-0
lines changed

‎src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
package com.thealgorithms.datastructures.trees;
2+
3+
public class LazySegmentTree {
4+
/**
5+
* Lazy Segment Tree
6+
*
7+
* @see
8+
* <a href="https://www.geeksforgeeks.org/lazy-propagation-in-segment-tree/">
9+
*/
10+
static class Node {
11+
private final int start, end; // start and end of the segment represented by this node
12+
private int value; // value is the sum of all elements in the range [start, end)
13+
private int lazy; // lazied value that should be added to children nodes
14+
Node left, right; // left and right children
15+
public Node(int start, int end, int value) {
16+
this.start = start;
17+
this.end = end;
18+
this.value = value;
19+
this.lazy = 0;
20+
this.left = null;
21+
this.right = null;
22+
}
23+
24+
/** Update the value of this node with the given value diff.
25+
*
26+
* @param diff The value to add to every index of this node range.
27+
*/
28+
public void applyUpdate(int diff) {
29+
this.lazy += diff;
30+
this.value += (this.end - this.start) * diff;
31+
}
32+
33+
/** Shift the lazy value of this node to its children.
34+
*/
35+
public void shift() {
36+
if (lazy == 0) return;
37+
if (this.left == null && this.right == null) return;
38+
this.value += this.lazy;
39+
if (this.left != null) this.left.applyUpdate(this.lazy);
40+
if (this.right != null) this.right.applyUpdate(this.lazy);
41+
this.lazy = 0;
42+
}
43+
44+
/** Create a new node that is the sum of this node and the given node.
45+
*
46+
* @param left The left Node of merging
47+
* @param right The right Node of merging
48+
* @return The new Node.
49+
*/
50+
static Node merge(Node left, Node right) {
51+
if (left == null) return right;
52+
if (right == null) return left;
53+
Node result = new Node(left.start, right.end, left.value + right.value);
54+
result.left = left;
55+
result.right = right;
56+
return result;
57+
}
58+
59+
public int getValue() {
60+
return value;
61+
}
62+
63+
public Node getLeft() {
64+
return left;
65+
}
66+
67+
public Node getRight() {
68+
return right;
69+
}
70+
}
71+
72+
private final Node root;
73+
74+
/** Create a new LazySegmentTree with the given array.
75+
*
76+
* @param array The array to create the LazySegmentTree from.
77+
*/
78+
public LazySegmentTree(int[] array) {
79+
this.root = buildTree(array, 0, array.length);
80+
}
81+
82+
/** Build a new LazySegmentTree from the given array in O(n) time.
83+
*
84+
* @param array The array to build the LazySegmentTree from.
85+
* @param start The start index of the current node.
86+
* @param end The end index of the current node.
87+
* @return The root of the new LazySegmentTree.
88+
*/
89+
private Node buildTree(int[] array, int start, int end) {
90+
if (end - start < 2) return new Node(start, end, array[start]);
91+
int mid = (start + end) >> 1;
92+
Node left = buildTree(array, start, mid);
93+
Node right = buildTree(array, mid, end);
94+
return Node.merge(left, right);
95+
}
96+
97+
/** Update the value of given range with the given value diff in O(log n) time.
98+
*
99+
* @param left The left index of the range to update.
100+
* @param right The right index of the range to update.
101+
* @param diff The value to add to every index of the range.
102+
* @param curr The current node.
103+
*/
104+
private void updateRange(int left, int right, int diff, Node curr) {
105+
if (left <= curr.start && curr.end <= right) {
106+
curr.applyUpdate(diff);
107+
return;
108+
}
109+
if (left >= curr.end || right <= curr.start) return;
110+
curr.shift();
111+
updateRange(left, right, diff, curr.left);
112+
updateRange(left, right, diff, curr.right);
113+
Node merge = Node.merge(curr.left, curr.right);
114+
curr.value = merge.value;
115+
}
116+
117+
/** Get Node of given range in O(log n) time.
118+
*
119+
* @param left The left index of the range to update.
120+
* @param right The right index of the range to update.
121+
* @return The Node representing the sum of the given range.
122+
*/
123+
private Node getRange(int left, int right, Node curr) {
124+
if (left <= curr.start && curr.end <= right) return curr;
125+
if (left >= curr.end || right <= curr.start) return null;
126+
curr.shift();
127+
return Node.merge(getRange(left, right, curr.left), getRange(left, right, curr.right));
128+
}
129+
130+
public int getRange(int left, int right) {
131+
Node result = getRange(left, right, root);
132+
return result == null ? 0 : result.getValue();
133+
}
134+
135+
public void updateRange(int left, int right, int diff) {
136+
updateRange(left, right, diff, root);
137+
}
138+
139+
public Node getRoot() {
140+
return root;
141+
}
142+
}

‎src/test/java/com/thealgorithms/datastructures/trees/LazySegmentTreeTest.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
package com.thealgorithms.datastructures.trees;
2+
3+
import org.junit.jupiter.api.Test;
4+
5+
import static org.junit.jupiter.api.Assertions.*;
6+
7+
public class LazySegmentTreeTest {
8+
9+
@Test
10+
void build() {
11+
int[] arr = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
12+
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
13+
assertEquals(55, lazySegmentTree.getRoot().getValue());
14+
assertEquals(15, lazySegmentTree.getRoot().getLeft().getValue());
15+
assertEquals(40, lazySegmentTree.getRoot().getRight().getValue());
16+
}
17+
18+
@Test
19+
void update() {
20+
int[] arr = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
21+
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
22+
assertEquals(10, lazySegmentTree.getRoot().getValue());
23+
24+
lazySegmentTree.updateRange(0, 2, 1);
25+
assertEquals(12, lazySegmentTree.getRoot().getValue());
26+
27+
lazySegmentTree.updateRange(1, 3, 1);
28+
assertEquals(14, lazySegmentTree.getRoot().getValue());
29+
30+
lazySegmentTree.updateRange(6, 8, 1);
31+
assertEquals(16, lazySegmentTree.getRoot().getValue());
32+
33+
lazySegmentTree.updateRange(3, 9, 1);
34+
assertEquals(22, lazySegmentTree.getRoot().getValue());
35+
}
36+
37+
@Test
38+
void get() {
39+
int[] arr = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
40+
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
41+
assertEquals(55, lazySegmentTree.getRange(0, 10));
42+
assertEquals(3, lazySegmentTree.getRange(0, 2));
43+
assertEquals(19, lazySegmentTree.getRange(8, 10));
44+
assertEquals(44, lazySegmentTree.getRange(1, 9));
45+
}
46+
47+
@Test
48+
void updateAndGet() {
49+
int[] arr = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
50+
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
51+
52+
for (int i = 0; i < 10; i++) for (int j = i + 1; j < 10; j++) {
53+
lazySegmentTree.updateRange(i, j, 1);
54+
assertEquals(j - i, lazySegmentTree.getRange(i, j));
55+
lazySegmentTree.updateRange(i, j, -1);
56+
assertEquals(0, lazySegmentTree.getRange(i, j));
57+
}
58+
}
59+
60+
}

0 commit comments

Comments
 (0)