Getting TLE in python


#1

I have used a stack to find out prev greater element and next greater element in O(n) time. I have also used an efficient implementation of power function still I am getting TLE. Below is my code:
Please help me figure out what I am doing wrong.

Blockquote
import math
class Solution:
# @param A : list of integers
# @param B : list of integers
# @return a list of integers

def pow(self, x, n, d):
    if x == 0:
        return 0
    if n == 0:
        return 1
    temp = self.pow(x,n//2,d)
    if n % 2 == 0:
        return ((temp%d) * (temp %d)) % d
    else :
        return ( (x % d) * (((temp%d) * (temp %d)) % d )) % d
    

def divisor(self,num):
    c = 0
    for i in range(1,int(math.sqrt(num))+1):
        if num % i == 0:
            if num / i == i:
                c += 1
            else :
                c += 2
    return c
 
 
 help               
def solve(self, A, B):
    x1 = [ 1 for i in range(len(A))]
    x2 = [ 1 for i in range(len(A))]
    stack = []
    
    
    # next greater element
    stack.append([A[0],0])
    n = len(A)
    for i in range(1,n):
        if stack[-1][0] > A[i]:
            stack.append([A[i],i])
        else :
            while(len(stack)!=0 and stack[-1][0] <= A[i]):
                x1[stack[-1][1]] = i - stack[-1][1]
                stack.pop(-1)
            stack.append([A[i],i])
                
    while(len(stack)!=0):
        x1[stack[-1][1]] = n - stack[-1][1]
        stack.pop(-1)
    
    # previous greater element
    stack.append([A[n-1],n-1])
    n = len(A)
    for i in reversed(range(0,n-1)):
        if stack[-1][0] >= A[i]:
            stack.append([A[i],i])
        else :
            while(len(stack)!=0 and stack[-1][0] < A[i]):
                x2[stack[-1][1]] = stack[-1][1] - i
                stack.pop(-1)
            stack.append([A[i],i])
                
    while(len(stack)!=0):
        x2[stack[-1][1]] = stack[-1][1] + 1
        stack.pop(-1)
        
    
    #divisor product for each element in A
    temp = {}
    for i in range(len(A)):
        if A[i] not in temp :
            div = self.divisor(A[i])
            if div % 2== 0:
                temp[A[i]] = self.pow(A[i],div//2,1000000007)
            else :
                temp[A[i]] = (int(math.sqrt(A[i]) * self.pow(A[i],div//2,1000000007))) % 1000000007
    
    
    #2D array [divisor product, frequency]
    a = [[0,0] for i in range(n)]
    for i  in range(n):
        a[i] = [temp[A[i]],x1[i] * x2[i]]
    
    a.sort(reverse= True, key = lambda x : (x[0],x[1]))
    
    
    #prfeix sum array of frequencies of divisor products
    y = [ 0 for i in range(n)]
    y[0] = a[0][1]
    for i in range(1,n):
        y[i] += y[i-1] + a[i][1]
    #print(y[n-1])
    
    
    #binary search to find result of each query
    res = []
    for i in range(len(B)):
        l = 0
        r = n - 1
        ans = 0
        while(l<=r):
            m = l + (r-l)//2
            if y[m] < B[i]:
                l = m + 1
            else :
                ans = m
                r = m - 1
        res.append(a[ans][0])
        
                
    return res

Blockquote


#2

I am getting the same.
In my case, it gives TLE in python3 but works fine with python2.