diff --git a/src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java b/src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java new file mode 100644 index 000000000..fe1bf8ec4 --- /dev/null +++ b/src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java @@ -0,0 +1,142 @@ +package com.thealgorithms.datastructures.trees; + +public class LazySegmentTree { + /** + * Lazy Segment Tree + * + * @see + * + */ + 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; + } +} diff --git a/src/test/java/com/thealgorithms/datastructures/trees/LazySegmentTreeTest.java b/src/test/java/com/thealgorithms/datastructures/trees/LazySegmentTreeTest.java new file mode 100644 index 000000000..69153c7d2 --- /dev/null +++ b/src/test/java/com/thealgorithms/datastructures/trees/LazySegmentTreeTest.java @@ -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)); + } + } + +}