Array range queries to find the number of perfect square elements with updates | Python program to find number of perfect square numbers in a subarray and performing updates ; Function to check if a number is a perfect square or not ; Find floating point value of square root of x . ; If square root is an integer ; A utility function to get the middle index from corner indexes . ; where st -- > Pointer to segment tree index -- > Index of current node in the segment tree . Initially 0 is passed as root is always at index 0 ss & se -- > Starting and ending indexes of the segment represented by current node i . e . st [ index ] qs & qe -- > Starting and ending indexes of query range ; If segment of this node is a part of given range , then return the number of perfect square numbers in the segment ; If segment of this node is outside the given range ; If a part of this segment overlaps with the given range ; where st , si , ss & se are same as getSumUtil ( ) i -- > index of the element to be updated . This index is in input array . diff -- > Value to be added to all nodes which have i in range ; Base Case : If the input index lies outside the range of this segment ; If the input index is in range of this node , then update the value of the node and its children ; Function to update a value in the input array and segment tree . It uses updateValueUtil ( ) to update the value in segment tree ; Check for erroneous input index ; Update the value in array ; Case 1 : Old and new values both are perfect square numbers ; Case 2 : Old and new values both not perfect square numbers ; Case 3 : Old value was perfect square , new value is not a perfect square ; Case 4 : Old value was non - perfect square , new_val is perfect square ; Update values of nodes in segment tree ; Return no . of perfect square numbers in range from index qs ( query start ) to qe ( query end ) . It mainly uses queryUtil ( ) ; Recursive function that constructs Segment Tree for array [ ss . . se ] . si is index of current node in segment tree st ; If there is one element in array , check if it is perfect square number then store 1 in the segment tree else store 0 and return ; if arr [ ss ] is a perfect square number ; If there are more than one elements , then recur for left and right subtrees and store the sum of the two values in this node ; Function to construct a segment tree from given array . This function allocates memory for segment tree and calls constructSTUtil ( ) to fill the allocated memory ; Allocate memory for segment tree Height of segment tree ; Maximum size of segment tree ; Fill the allocated memory st ; Return the constructed segment tree ; Driver Code ; Build segment tree from given array ; Query 1 : Query ( start = 0 , end = 4 ) ; Query 2 : Update ( i = 3 , x = 11 ) , i . e Update a [ i ] to x ; Query 3 : Query ( start = 0 , end = 4 )
from math import sqrt , floor , ceil , log2 NEW_LINE from typing import List NEW_LINE MAX = 1000 NEW_LINE def isPerfectSquare ( x : int ) -> bool : NEW_LINE INDENT sr = sqrt ( x ) NEW_LINE return True if ( ( sr - floor ( sr ) ) == 0 ) else False NEW_LINE DEDENT def getMid ( s : int , e : int ) -> int : NEW_LINE INDENT return s + ( e - s ) // 2 NEW_LINE DEDENT def queryUtil ( st : List [ int ] , ss : int , se : int , qs : int , qe : int , index : int ) -> int : NEW_LINE INDENT if ( qs <= ss and qe >= se ) : NEW_LINE INDENT return st [ index ] NEW_LINE DEDENT if ( se < qs or ss > qe ) : NEW_LINE INDENT return 0 NEW_LINE DEDENT mid = getMid ( ss , se ) NEW_LINE return queryUtil ( st , ss , mid , qs , qe , 2 * index + 1 ) + queryUtil ( st , mid + 1 , se , qs , qe , 2 * index + 2 ) NEW_LINE DEDENT def updateValueUtil ( st : List [ int ] , ss : int , se : int , i : int , diff : int , si : int ) -> None : NEW_LINE INDENT if ( i < ss or i > se ) : NEW_LINE INDENT return NEW_LINE DEDENT st [ si ] = st [ si ] + diff NEW_LINE if ( se != ss ) : NEW_LINE INDENT mid = getMid ( ss , se ) NEW_LINE updateValueUtil ( st , ss , mid , i , diff , 2 * si + 1 ) NEW_LINE updateValueUtil ( st , mid + 1 , se , i , diff , 2 * si + 2 ) NEW_LINE DEDENT DEDENT def updateValue ( arr : List [ int ] , st : List [ int ] , n : int , i : int , new_val : int ) -> None : NEW_LINE INDENT if ( i < 0 or i > n - 1 ) : NEW_LINE INDENT print ( " Invalid ▁ Input " ) NEW_LINE return NEW_LINE DEDENT diff = 0 NEW_LINE oldValue = 0 NEW_LINE oldValue = arr [ i ] NEW_LINE arr [ i ] = new_val NEW_LINE if ( isPerfectSquare ( oldValue ) and isPerfectSquare ( new_val ) ) : NEW_LINE INDENT return NEW_LINE DEDENT if ( not isPerfectSquare ( oldValue ) and not isPerfectSquare ( new_val ) ) : NEW_LINE INDENT return NEW_LINE DEDENT if ( isPerfectSquare ( oldValue ) and not isPerfectSquare ( new_val ) ) : NEW_LINE INDENT diff = - 1 NEW_LINE DEDENT if ( not isPerfectSquare ( oldValue ) and not isPerfectSquare ( new_val ) ) : NEW_LINE INDENT diff = 1 NEW_LINE DEDENT updateValueUtil ( st , 0 , n - 1 , i , diff , 0 ) NEW_LINE DEDENT def query ( st : List [ int ] , n : int , qs : int , qe : int ) -> None : NEW_LINE INDENT perfectSquareInRange = queryUtil ( st , 0 , n - 1 , qs , qe , 0 ) NEW_LINE print ( perfectSquareInRange ) NEW_LINE DEDENT def constructSTUtil ( arr : List [ int ] , ss : int , se : int , st : List [ int ] , si : int ) -> int : NEW_LINE INDENT if ( ss == se ) : NEW_LINE INDENT if ( isPerfectSquare ( arr [ ss ] ) ) : NEW_LINE INDENT st [ si ] = 1 NEW_LINE DEDENT else : NEW_LINE INDENT st [ si ] = 0 NEW_LINE DEDENT return st [ si ] NEW_LINE DEDENT mid = getMid ( ss , se ) NEW_LINE st [ si ] = constructSTUtil ( arr , ss , mid , st , si * 2 + 1 ) + constructSTUtil ( arr , mid + 1 , se , st , si * 2 + 2 ) NEW_LINE return st [ si ] NEW_LINE DEDENT def constructST ( arr : List [ int ] , n : int ) -> List [ int ] : NEW_LINE INDENT x = ( ceil ( log2 ( n ) ) ) NEW_LINE max_size = 2 * pow ( 2 , x ) - 1 NEW_LINE st = [ 0 for _ in range ( max_size ) ] NEW_LINE constructSTUtil ( arr , 0 , n - 1 , st , 0 ) NEW_LINE return st NEW_LINE DEDENT if __name__ == " _ _ main _ _ " : NEW_LINE INDENT arr = [ 16 , 15 , 8 , 9 , 14 , 25 ] NEW_LINE n = len ( arr ) NEW_LINE st = constructST ( arr , n ) NEW_LINE start = 0 NEW_LINE end = 4 NEW_LINE query ( st , n , start , end ) NEW_LINE i = 3 NEW_LINE x = 11 NEW_LINE updateValue ( arr , st , n , i , x ) NEW_LINE start = 0 NEW_LINE end = 4 NEW_LINE query ( st , n , start , end ) NEW_LINE DEDENT