Algorithms/JavaAlgorithms
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Lazy Segment Tree (#3209)

Browse Source
This commit is contained in:
AmirMohammad Hosseini Nasab 2022-08-10 21:57:53 +04:30 committed by GitHub
parent d63813e0a2
commit 3e8f30c72a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 202 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

142
src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java Normal file
View File

@ -0,0 +1,142 @@
package com.thealgorithms.datastructures.trees;
public class LazySegmentTree {
/**
* Lazy Segment Tree
*
* @see
* <a href="https://www.geeksforgeeks.org/lazy-propagation-in-segment-tree/">
*/
static class Node {
private final int start, end; // start and end of the segment represented by this node
private int value; // value is the sum of all elements in the range [start, end)
private int lazy; // lazied value that should be added to children nodes
Node left, right; // left and right children
public Node(int start, int end, int value) {
this.start = start;
this.end = end;
this.value = value;
this.lazy = 0;
this.left = null;
this.right = null;
}
/** Update the value of this node with the given value diff.
*
* @param diff The value to add to every index of this node range.
*/
public void applyUpdate(int diff) {
this.lazy += diff;
this.value += (this.end - this.start) * diff;
}
/** Shift the lazy value of this node to its children.
*/
public void shift() {
if (lazy == 0) return;
if (this.left == null && this.right == null) return;
this.value += this.lazy;
if (this.left != null) this.left.applyUpdate(this.lazy);
if (this.right != null) this.right.applyUpdate(this.lazy);
this.lazy = 0;
}
/** Create a new node that is the sum of this node and the given node.
*
* @param left The left Node of merging
* @param right The right Node of merging
* @return The new Node.
*/
static Node merge(Node left, Node right) {
if (left == null) return right;
if (right == null) return left;
Node result = new Node(left.start, right.end, left.value + right.value);
result.left = left;
result.right = right;
return result;
}
public int getValue() {
return value;
}
public Node getLeft() {
return left;
}
public Node getRight() {
return right;
}
}
private final Node root;
/** Create a new LazySegmentTree with the given array.
*
* @param array The array to create the LazySegmentTree from.
*/
public LazySegmentTree(int[] array) {
this.root = buildTree(array, 0, array.length);
}
/** Build a new LazySegmentTree from the given array in O(n) time.
*
* @param array The array to build the LazySegmentTree from.
* @param start The start index of the current node.
* @param end The end index of the current node.
* @return The root of the new LazySegmentTree.
*/
private Node buildTree(int[] array, int start, int end) {
if (end - start < 2) return new Node(start, end, array[start]);
int mid = (start + end) >> 1;
Node left = buildTree(array, start, mid);
Node right = buildTree(array, mid, end);
return Node.merge(left, right);
}
/** Update the value of given range with the given value diff in O(log n) time.
*
* @param left The left index of the range to update.
* @param right The right index of the range to update.
* @param diff The value to add to every index of the range.
* @param curr The current node.
*/
private void updateRange(int left, int right, int diff, Node curr) {
if (left <= curr.start && curr.end <= right) {
curr.applyUpdate(diff);
return;
}
if (left >= curr.end || right <= curr.start) return;
curr.shift();
updateRange(left, right, diff, curr.left);
updateRange(left, right, diff, curr.right);
Node merge = Node.merge(curr.left, curr.right);
curr.value = merge.value;
}
/** Get Node of given range in O(log n) time.
*
* @param left The left index of the range to update.
* @param right The right index of the range to update.
* @return The Node representing the sum of the given range.
*/
private Node getRange(int left, int right, Node curr) {
if (left <= curr.start && curr.end <= right) return curr;
if (left >= curr.end || right <= curr.start) return null;
curr.shift();
return Node.merge(getRange(left, right, curr.left), getRange(left, right, curr.right));
}
public int getRange(int left, int right) {
Node result = getRange(left, right, root);
return result == null ? 0 : result.getValue();
}
public void updateRange(int left, int right, int diff) {
updateRange(left, right, diff, root);
}
public Node getRoot() {
return root;
}
}

60
src/test/java/com/thealgorithms/datastructures/trees/LazySegmentTreeTest.java Normal file
View File

@ -0,0 +1,60 @@
package com.thealgorithms.datastructures.trees;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;
public class LazySegmentTreeTest {
@Test
void build() {
int[] arr = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
assertEquals(55, lazySegmentTree.getRoot().getValue());
assertEquals(15, lazySegmentTree.getRoot().getLeft().getValue());
assertEquals(40, lazySegmentTree.getRoot().getRight().getValue());
}
@Test
void update() {
int[] arr = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
assertEquals(10, lazySegmentTree.getRoot().getValue());
lazySegmentTree.updateRange(0, 2, 1);
assertEquals(12, lazySegmentTree.getRoot().getValue());
lazySegmentTree.updateRange(1, 3, 1);
assertEquals(14, lazySegmentTree.getRoot().getValue());
lazySegmentTree.updateRange(6, 8, 1);
assertEquals(16, lazySegmentTree.getRoot().getValue());
lazySegmentTree.updateRange(3, 9, 1);
assertEquals(22, lazySegmentTree.getRoot().getValue());
}
@Test
void get() {
int[] arr = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
assertEquals(55, lazySegmentTree.getRange(0, 10));
assertEquals(3, lazySegmentTree.getRange(0, 2));
assertEquals(19, lazySegmentTree.getRange(8, 10));
assertEquals(44, lazySegmentTree.getRange(1, 9));
}
@Test
void updateAndGet() {
int[] arr = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
LazySegmentTree lazySegmentTree = new LazySegmentTree(arr);
for (int i = 0; i < 10; i++) for (int j = i + 1; j < 10; j++) {
lazySegmentTree.updateRange(i, j, 1);
assertEquals(j - i, lazySegmentTree.getRange(i, j));
lazySegmentTree.updateRange(i, j, -1);
assertEquals(0, lazySegmentTree.getRange(i, j));
}
}
}