The Algorithms logo
The Algorithms
AboutDonate

Segment Tree Multiplication

using System;

namespace DataStructures.SegmentTrees
{
    /// <summary>
    ///     This is an extension of a segment tree, which allows applying distributive operations to a subarray
    ///     (in this case multiplication).
    /// </summary>
    public class SegmentTreeApply : SegmentTree
    {
        /// <summary>
        ///     Initializes a new instance of the <see cref="SegmentTreeApply" /> class.
        ///     Runtime complexity: O(n) where n equals the array-length.
        /// </summary>
        /// <param name="arr">Array on which the operations should be made.</param>
        public SegmentTreeApply(int[] arr)
            : base(arr)
        {
            // Initilizes and fills "operand" array with neutral element (in this case 1, because value * 1 = value)
            Operand = new int[Tree.Length];
            Array.Fill(Operand, 1);
        }

        /// <summary>
        ///     Gets an array that stores for each node an operand,
        ///     which must be applied to all direct and indirect child nodes of this node
        ///     (but not to the node itself).
        /// </summary>
        public int[] Operand { get; }

        /// <summary>
        ///     Applies a distributive operation to a subarray defined by <c>l</c> and <c>r</c>
        ///     (in this case multiplication by <c>value</c>).
        ///     Runtime complexity: O(logN) where N equals the initial array-length.
        /// </summary>
        /// <param name="l">Left border of the subarray.</param>
        /// <param name="r">Right border of the subarray.</param>
        /// <param name="value">Value with which each element of the interval is calculated.</param>
        public void Apply(int l, int r, int value)
        {
            // The Application start at node with 1
            // Node with index 1 includes the whole input subarray
            Apply(++l, ++r, value, 1, Tree.Length / 2, 1);
        }

        /// <summary>
        ///     Edits a query.
        /// </summary>
        /// <param name="l">Left border of the query.</param>
        /// <param name="r">Right border of the query.</param>
        /// <param name="a">Left end of the subarray enclosed by <c>i</c>.</param>
        /// <param name="b">Right end of the subarray enclosed by <c>i</c>.</param>
        /// <param name="i">Current node.</param>
        /// <returns>Sum of a subarray between <c>l</c> and <c>r</c> (including <c>l</c> and <c>r</c>).</returns>
        protected override int Query(int l, int r, int a, int b, int i)
        {
            if (l <= a && b <= r)
            {
                return Tree[i];
            }

            if (r < a || b < l)
            {
                return 0;
            }

            var m = (a + b) / 2;

            // Application of the saved operand to the direct and indrect child nodes
            return Operand[i] * (Query(l, r, a, m, Left(i)) + Query(l, r, m + 1, b, Right(i)));
        }

        /// <summary>
        ///     Applies the operation.
        /// </summary>
        /// <param name="l">Left border of the Application.</param>
        /// <param name="r">Right border of the Application.</param>
        /// <param name="value">Multiplier by which the subarray is to be multiplied.</param>
        /// <param name="a">Left end of the subarray enclosed by <c>i</c>.</param>
        /// <param name="b">Right end of the subarray enclosed by <c>i</c>.</param>
        /// <param name="i">Current node.</param>
        private void Apply(int l, int r, int value, int a, int b, int i)
        {
            // If a and b are in the (by l and r) specified subarray
            if (l <= a && b <= r)
            {
                // Applies the operation to the current node and saves it for the direct and indirect child nodes
                Operand[i] = value * Operand[i];
                Tree[i] = value * Tree[i];
                return;
            }

            // If a or b are out of the by l and r specified subarray stop application at this node
            if (r < a || b < l)
            {
                return;
            }

            // Calculates index m of the node that cuts the current subarray in half
            var m = (a + b) / 2;

            // Applies the operation to both halfes
            Apply(l, r, value, a, m, Left(i));
            Apply(l, r, value, m + 1, b, Right(i));

            // Recalculates the value of this node by its (possibly new) children.
            Tree[i] = Operand[i] * (Tree[Left(i)] + Tree[Right(i)]);
        }
    }
}