# Codeforces problem python code optimization

I have solved one problem on Codeforces and I guess the output is quite correct. However, it is slow for input in range of thousands. Since, I am learning Python now, it would be great help if one could help me optimise this code. Any notes, tips or suggestions on the sidelines are always welcome.

The question is a practice question and hence does not violate any contest rules.

import sys
import time

_no_of_ants, _no_of_queries = (input().split())
no_of_ants = int(_no_of_ants)
no_of_queries = int(_no_of_queries)

class ants_cord():
def __init__(self,point,no_of_ants):
self.co_ant = {point:no_of_ants}

def ants_at_point(self,point):
return self.co_ant[point]

def move_ants(self,point):
x1,y1 = point
if (x1+1,y1) in self.co_ant.keys():
self.co_ant[((x1+1),y1)] += 1
else:
self.co_ant[((x1+1),y1)] = 1
if (x1,y1+1) in self.co_ant.keys():
self.co_ant[((x1),y1+1)] += 1
else:
self.co_ant[((x1),y1+1)] = 1
if (x1,y1-1) in self.co_ant.keys():
self.co_ant[((x1),y1-1)] += 1
else:
self.co_ant[((x1),y1-1)] = 1
if (x1-1,y1) in self.co_ant.keys():
self.co_ant[((x1-1),y1)] += 1
else:
self.co_ant[((x1-1),y1)] = 1
self.co_ant[point] -= 4

def check_status(self):
for key in self.co_ant.keys():
if k.co_ant[key] >= 4:
return key
else: return False
start = time.time()
k = ants_cord((0,0),no_of_ants)
point = k.check_status()
while point:
k.move_ants(point)
point = k.check_status()
end = time.time()
print("Sec:{}".format(end-start))
while no_of_queries:
_x,_y = input().split()
x = int(_x)
y = int(_y)
if (x,y) in k.co_ant:
print(k.co_ant[(x,y)])
else:
print('0')
no_of_queries -= 1


EDIT: I updated the code according to suggestion. But I still get a TLE error in test case 4. Should I employ other algorithm then?

• I've added an answer. I don't have the right version of python handy but it should again make your script orders of magnitude faster. I'd be interested in the results if you were to perform performances tests with a high number of ants. – SylvainD Jun 24 '13 at 14:51
• I'm not sure that the task is solvable using python or any interpreted language at all. As you can see there aren't any python/ruby/.. solution and the fastest c++ solution time is 100 ms, so it's just impossible to run faster than may be ~2-5 seconds for a python solution. There is the list of failed (due to time limit) solutions - just apply the status filter (Python 2 or Python 3 and the Verdict like the 'Time limit...'). – cat_baxter Jun 24 '13 at 15:29

Here are a few improvements you could perform on the algorithm itself :

• You could simulate many iterations in one go.

In order to do so,

def move_ants(g,point):
x,y = point
for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
g[a,b] = 1 + g.get((a,b),0)
g[point] -= 4


could be changed for :

def move_ants(g,point):
x,y = point
nb_ants = g[point]
nb_ants_per_dir = nb_ants/4
for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
g[a,b] = nb_ants_per_dir + g.get((a,b),0)
g[point] = nb_ants - 4 * nb_ants_per_dir

• check_status could return a set of points to be handled

You'd just need to do something like :

def check_status(g):
return [k for k,v in g.iteritems() if v >= 4]

• Also, you might want to keep track of the cells you've updated in order no to loop over the all dict but only the candidates. This would need quite a lot of rewritting.

• Also, you could take into account the fact that the problem offer many axis of symetry : obviously the x and y axis but also the y=x and the y=-x. Thus you could restrict yourself to 1/8 of the plan corresponding to 0 < x < y

Edit

After testing, this code on 5555 ants:

  def move_ants(g,points):
cand=set()
for point in points:
x,y = point
nb_ants = g[point]
nb_ants_per_dir = nb_ants/4
for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
new_nb = g.get((a,b),0) + nb_ants_per_dir
g[a,b] = new_nb
if new_nb >= 4:
g[point] = nb_ants - 4*nb_ants_per_dir
return cand

g = {(0,0) : no_of_ants} # or you could use default dict
points = [(0,0)]
while points:
points = move_ants(g,points)
print("Sec:{}".format(time.time()-start))
for q in xrange(no_of_queries):
x,y=1,1
print(g.get((x,y),0))


appears to be 21 times faster than the solution provided in my other answer and 141 times faster than yours.

A few things:

• You should probably put the start = time.time() as early as possible if you want it to make sense.

• As your class consists on a simple dict, there is no or little value in defining it. (In a "real life" scenario, it would be considered as a good practice).

Things got 25% quicker on average starting with 1000 ants.

• You could make your code more concise by using structure to isolate the parts that change.

move_ants() would look like this :

def move_ants(g,point):
x,y = point
for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
if (a,b) in g.keys():
g[a,b] += 1
else:
g[a,b] = 1
g[point] -= 4


(Not performance impact)

• In check_status(), instead of iterating on keys and then retrieve the corresponding value, you could use iteritems() :

check_status() would be like this :

def check_status(g):
for k,v in g.iteritems():
if v >= 4:
return k
else: return False


(This does improve performances)

• To get the values from your dict with default values, you could use get(). You'd have something like print(g.get((x,y),0)) and g[a,b] = 1 + g.get((a,b),0). (I really didn't expect so but this made the script more than twice faster).

At the end, my script is like this (I've changed other details but it's mostly personal preferences) :

#!/usr/bin/python

import sys
import time

start = time.time()

##_no_of_ants, _no_of_queries = (input().split())
##no_of_ants = int(_no_of_ants)
##no_of_queries = int(_no_of_queries)

no_of_ants = 1000
no_of_queries = 1

def move_ants(g,point):
x,y = point
for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
g[a,b] = 1 + g.get((a,b),0)
g[point] -= 4

def check_status(g):
for k,v in g.iteritems():
if v >= 4:
return k
else: return False

g = {(0,0) : no_of_ants} # or you could use default dict
point = check_status(g)
while point:
move_ants(g,point)
point = check_status(g)
print("Sec:{}".format(time.time()-start))
for q in xrange(no_of_queries):
###_x,_y = input().split()
###x = int(_x)
###y = int(_y)
x,y=1,1
print(g.get((x,y),0))

• It would worthy to mention that your code is in Python 2.x. Thanks for response :) – Vivek Rai Jun 23 '13 at 19:49
• The g[a,b] = 1 + g.get((a,b),0) is already nice, but still a double dict access. The prototypical tip here is to benchmark also using either setdefault() or collections.defaultdict(). – user2512319 Jun 24 '13 at 4:53
• user2512319 : setdefault is what I wanted to do as well but I couldn't make it work with integers. As for defaultdict i considered it as an option (cf comment in the code) and then forgot about it. – SylvainD Jun 24 '13 at 7:11