To count number of quadruplets with sum = 0

[Michael Spencer]

n00m wrote:
> http://www.spoj.pl/problems/SUMFOUR/
> 
> 3
> 0 0 0 0
> 0 0 0 0
> -1 -1 1 1
> Answer for this input data is 33.
> 
> My solution for the problem is
> ======================================================================
> 
> import time
> t = time.clock()
> 
> q,w,e,r,sch,h = [],[],[],[],0,{}
> 
> f = open("D:/m4000.txt","rt")
> 
> n = int(f.readline())
> 
> for o in range(n):
>   row = map(long, f.readline().split())
>   q.append(row[0])
>   w.append(row[1])
>   e.append(row[2])
>   r.append(row[3])
> 
> f.close()
> 
> for x in q:
>   for y in w:
>     if h.has_key(x+y):
>       h[x+y] += 1
>     else:
>       h[x+y] = 1
> 
> for x in e:
>   for y in r:
>     sch += h.get(-(x+y),0)
> 
> q,w,e,r,h = None,None,None,None,None
> 
> print sch
> print time.clock() - t
> 
> ===============================================================
> 
> Alas it gets "time limit exceeded".
> On my home PC (1.6 GHz, 512 MB RAM) and for 4000 input rows it
> executes ~1.5 min.
> Any ideas to speed it up say 10 times? Or the problem only for C-like
> langs?
> 
Perhaps a bit faster using slicing to get the lists and avoiding dict.get:

def sumfour(src):
     l = map(int, src.split())
     dct={}
     s=0
     A, B, C, D = l[1::4], l[2::4], l[3::4], l[4::4]
     for a in A:
         for b in B:
             if a+b in dct:
                 dct[a+b] += 1
             else:
                 dct[a+b] = 1
     for c in C:
         for d in D:
             if -c-d in dct:
                 s+= dct[-c-d]
     return s


if __name__ == '__main__':
     import sys
     print sumfour(sys.stdin.read())


Michael