p229

p229.py

def segment(limit):
    
    segmentstart = 0
    segmentsize = limit//1000 #adjust according to limit
    
    b7max = [0]*(int(math.sqrt(limit)) + 1) #Keeps track of maximum b given a for each segment
    b3max = [0]*(int(math.sqrt(limit)) + 1)
    b2max = [0]*(int(math.sqrt(limit)) + 1)
    b1max = [0]*(int(math.sqrt(limit)) + 1)
    
    total = 0
    while segmentstart <= limit:
        print(segmentstart)
        segmentend = min(segmentstart + segmentsize, limit)
        
        sqs7 = set([])
        sqs3 = set([])
        sqs2 = set([])
        sqs1 = set([])
        
        for a in range(1, int(math.sqrt(segmentend)) + 1):
            for b in range(b7max[a] + 1, int(math.sqrt((segmentend - a*a)//7)) + 1):
                v = a*a + 7*b*b
                sqs7.add(v)
            b7max[a] = b
                
        for a in range(1, int(math.sqrt(segmentend)) + 1):
            for b in range(b3max[a] + 1, int(math.sqrt((segmentend - a*a)//3)) + 1):
                v = a*a + 3*b*b
                if v in sqs7:
                    sqs3.add(v)
            b3max[a] = b
        
        for a in range(1, int(math.sqrt(segmentend)) + 1):
            for b in range(b2max[a] + 1, int(math.sqrt((segmentend - a*a)//2)) + 1):
                v = a*a + 2*b*b
                if v in sqs3:
                    sqs2.add(v)
            b2max[a] = b    
            
        for a in range(1, int(math.sqrt(segmentend)) + 1):
            for b in range(b1max[a] + 1, int(math.sqrt((segmentend - a*a))) + 1):
                v = a*a + b*b
                if v in sqs2:
                    sqs1.add(v)
            b1max[a] = b
            
        segmentstart += segmentsize
        total += len(sqs1)
    return total