Question

I’ve written this very badly optimized C code that does a simple math calculation:

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
#define MAX(a, b) (((a) > (b)) ? (a) : (b))


unsigned long long int p(int);
float fullCheck(int);

int main(int argc, char **argv){
  int i, g, maxNumber;
  unsigned long long int diff = 1000;

  if(argc < 2){
    fprintf(stderr, "Usage: %s maxNumber\n", argv[0]);
    return 0;
  }
  maxNumber = atoi(argv[1]);

  for(i = 1; i < maxNumber; i++){
    for(g = 1; g < maxNumber; g++){
      if(i == g)
        continue;
      if(p(MAX(i,g)) - p(MIN(i,g)) < diff &&  fullCheck(p(MAX(i,g)) - p(MIN(i,g))) && fullCheck(p(i) + p(g))){
          diff = p(MAX(i,g)) - p(MIN(i,g));
          printf("We have a couple %llu %llu with diff %llu\n", p(i), p(g), diff);
      }
    }
  }

  return 0;
}

float fullCheck(int number){
  float check = (-1 + sqrt(1 + 24 * number))/-6;
  float check2 = (-1 - sqrt(1 + 24 * number))/-6;
  if(check/1.00 == (int)check)
    return check;
  if(check2/1.00 == (int)check2)
    return check2;
  return 0;
}

unsigned long long int p(int n){
  return n * (3 * n - 1 ) / 2;
}

And then I’ve tried (just for fun) to port it under Python to see how it would react. My first version was almost a 1:1 conversion that run terribly slow (120+secs in Python vs <1sec in C).
I’ve done a bit of optimization, and this is what I obtained:

#!/usr/bin/env/python
from cmath import sqrt
import cProfile
from pstats import Stats

def quickCheck(n):
        partial_c = (sqrt(1 + 24 * (n)))/-6 
        c = 1/6 + partial_c
        if int(c.real) == c.real:
                return True
        c = c - 2*partial_c
        if int(c.real) == c.real:
                return True
        return False

def main():        
        maxNumber = 5000
        diff = 1000
        for i in range(1, maxNumber):
                p_i = i * (3 * i - 1 ) / 2
                for g in range(i, maxNumber):
                        if i == g:
                                continue
                        p_g = g * (3 * g - 1 ) / 2
                        if p_i > p_g:
                                ma = p_i
                                mi = p_g
                        else:
                                ma = p_g
                                mi = p_i

                        if ma - mi < diff and quickCheck(ma - mi):
                                if quickCheck(ma + mi):
                                        print ('New couple ', ma, mi)
                                        diff = ma - mi


cProfile.run('main()','script_perf')
perf = Stats('script_perf').sort_stats('time', 'calls').print_stats(10)

This runs in about 16secs which is better but also almost 20 times slower than C.
Now, I know C is better than Python for this kind of calculations, but what I would like to know is if there something that I’ve missed (Python-wise, like an horribly slow function or such) that could have made this function faster.
Please note that I’m using Python 3.1.1, if this makes a difference

Answer 1

I made it go from ~7 seconds to ~3 seconds on my machine:

• Precomputed i * (3 * i - 1 ) / 2 for each value, in yours it was computed twice quite a lot
• Cached calls to quickCheck
• Removed if i == g by adding +1 to the range
• Removed if p_i > p_g since p_i is always smaller than p_g

Also put the quickCheck-function inside main, to make all variables local (which have faster lookup than global).
I’m sure there are more micro-optimizations available.

def main():
        maxNumber = 5000
        diff = 1000

        p = {}
        quickCache = {}

        for i in range(maxNumber):
            p[i] = i * (3 * i - 1 ) / 2

        def quickCheck(n):
            if n in quickCache: return quickCache[n]
            partial_c = (sqrt(1 + 24 * (n)))/-6 
            c = 1/6 + partial_c
            if int(c.real) == c.real:
                    quickCache[n] = True
                    return True
            c = c - 2*partial_c
            if int(c.real) == c.real:
                    quickCache[n] = True
                    return True
            quickCache[n] = False
            return False

        for i in range(1, maxNumber):
                mi = p[i]
                for g in range(i+1, maxNumber):
                        ma = p[g]
                        if ma - mi < diff and quickCheck(ma - mi) and quickCheck(ma + mi):
                                print('New couple ', ma, mi)
                                diff = ma - mi