src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java - master - NikolayIvkin/TheAlgorithms

src/main/java/com/thealgorithms/datastructures/trees/LazySegmentTree.java
174 строки5 KB
Samuel Facchinello
style: enable `NeedBraces` in checkstyle (#5227)
13 июн 2024, 22:00
13 июн 2024, 22:0087b17e0
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;
        private final int 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;
        Node right; // left and right children

        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;
    }
}