Advertisement
Guest User

Untitled

a guest
Oct 31st, 2014
143
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 40.58 KB | None | 0 0
  1. import re
  2. #from math import sqrt
  3. from sys import stdin
  4. from array import array
  5. from math import sqrt
  6.  
  7. INTEGER_RE = "(\d+)" #Matches any non negative integer
  8. DOUBLE_RE = "(\d+\.\d*)" #INVARIANT: x,y positive => reg exp supporting negative floating points "([-]?\d+\.\d*)" not needed
  9. SEPARATOR = ' ' #We consider just spaces as separators
  10. #Regular expression for a topic
  11. TOPIC_REGEXP = INTEGER_RE + SEPARATOR + DOUBLE_RE + SEPARATOR + DOUBLE_RE + SEPARATOR + '*'
  12. #Regular expression for a query
  13. QUERY_REGEXP = '([a-zA-Z]{1})' + SEPARATOR + INTEGER_RE + SEPARATOR + DOUBLE_RE + SEPARATOR + DOUBLE_RE + SEPARATOR + '*'
  14.  
  15. #DEBUG
  16. #file_out = open('nearby_results.txt', 'w')
  17. #file_log = open('nearby_debug.txt', 'w')
  18.  
  19. '''Reads the input from a file f
  20. The input is assumed to be formatted as follows:
  21. First line: 3 integers T Q N
  22. T lines composed by an integer and 2 doubles
  23. Q lines composed by 2 integers q_id Qn and then another Qn integers
  24. N lines composed by 1 char, 1 int and 2 doubles
  25. @param f: The file from which the input should be read;
  26. @return: (topics, questions, queries)
  27. Three lists containing the topics, questions and queries read from the input channel.
  28. '''
  29. def read_and_process_input(f):
  30. line = f.readline()
  31.  
  32. regex = re.compile(INTEGER_RE) #Regular Expression for integers
  33. #INVARIANT: the input is assumed well formed and adherent to the specs above
  34. [T,Q,N] = regex.findall(line)
  35.  
  36. T = int(T)
  37. Q = int(Q)
  38. N = int(N)
  39.  
  40. topics = {}
  41. questions = []
  42.  
  43. #Maximum number of points or children for each topic
  44. max_elements_per_cluster = 16
  45. split_size = max_elements_per_cluster / 2
  46.  
  47.  
  48. ''' Creates the empty root of an SS-tree.
  49. @return: The newly created tree root.
  50. '''
  51. def ss_make_tree():
  52. tree = {'points': [], 'leaf':True, 'parent': None, 'x':0., 'y':0.}
  53. return tree
  54.  
  55.  
  56. ''' Inserts a point (topic) in a SS-tree; If necessary, splits the tree node in which the point was inserted
  57. and fixes the tree structure from that node up to the root.
  58.  
  59. @param new_point: The point to be inserted;
  60. The point must be a dictionary with three fields:
  61. - 'x': Point's x coordinate;
  62. - 'y': Point's y coordinate;
  63. - 't_id': The topic's ID.
  64. @param tree: The root of the tree in which the point is going to be inserted;
  65. @return tree: The root of the tree, possibly a new one if the old root has been split.
  66. '''
  67. def ss_tree_insert(new_point, tree):
  68. x_new_point = new_point['x']
  69. y_new_point = new_point['y']
  70.  
  71. #Looks for the right leaf (the one with the closest centroid) to which the new_point should be added.
  72. #INVARIANT: The empty tree's root is a (empty) leaf.
  73. node = tree
  74. while not node['leaf']:
  75. children = node['children']
  76. child = children[0]
  77. min_dist = (child['x'] - x_new_point) ** 2 + (child['y'] - y_new_point) ** 2
  78. min_index = 0
  79. for i in range(1,len(children)):
  80. child = children[i]
  81. dist = (child['x'] - x_new_point) ** 2 + (child['y'] - y_new_point) ** 2
  82. if dist < min_dist:
  83. min_index = i
  84. min_dist = dist
  85. node = children[min_index]
  86.  
  87.  
  88. #Now adds the new point to the leaf it has found.
  89.  
  90. #INVARIANT: node is a leaf
  91. points = node['points']
  92. if len(points) < max_elements_per_cluster:
  93. #No split neeeded to add the point to this node
  94.  
  95. #Can add the new_point to this node
  96. old_x_node = x_node = node['x']
  97. old_y_node = y_node = node['y']
  98.  
  99. #Compute the new centroid for the node
  100. n_p = len(points)
  101. x_node *= n_p
  102. y_node *= n_p
  103. x_node += x_new_point
  104. y_node += y_new_point
  105. points.append(new_point)
  106. n_p += 1
  107. x_node /= n_p
  108. y_node /= n_p
  109. node['x'] = x_node
  110. node['y'] = y_node
  111.  
  112. #Compute node's radius and variance
  113. radius = 0.
  114. x_var = y_var = 0.
  115. for point in points:
  116. #INVARIANT: points don't have radius
  117. x_dist = (x_node - point['x']) ** 2
  118. y_dist = (y_node - point['y']) ** 2
  119. radius = max(radius, x_dist + y_dist)
  120. #We don't need the exact variance, we can do fine with an estimate based on max distance form the centroid
  121. x_var = max(x_var, x_dist)
  122. y_var = max(y_var, y_dist)
  123. node['radius'] = sqrt(radius)
  124. node['x_var'] = x_var
  125. node['y_var'] = y_var
  126.  
  127. #Propagates the change all the way to the root
  128. node_parent = node['parent']
  129. while node_parent != None:
  130. tmp_x = x_node_parent = node_parent['x']
  131. tmp_y = y_node_parent = node_parent['y']
  132. n_p = len(node_parent['children'])
  133. x_node_parent *= n_p
  134. y_node_parent *= n_p
  135. x_node_parent += x_node - old_x_node
  136. y_node_parent += y_node - old_y_node
  137. old_x_node = tmp_x
  138. old_y_node = tmp_y
  139. x_node_parent /= n_p
  140. y_node_parent /= n_p
  141. node_parent['x'] = x_node_parent
  142. node_parent['y'] = y_node_parent
  143.  
  144. radius = 0.
  145. x_var = y_var = 0.
  146. for child_node in node_parent['children']:
  147. x_dist = (x_node_parent - child_node['x']) ** 2
  148. y_dist = (y_node_parent - child_node['y']) ** 2
  149. radius = max(radius, sqrt(x_dist + y_dist) + child_node['radius'])
  150. #We don't need the exact variance, we can do fine with an estimate based on max distance form the centroid
  151. x_var = max(x_var, x_dist + child_node['radius'] ** 2)
  152. y_var = max(y_var, y_dist + child_node['radius'] ** 2)
  153.  
  154. node_parent['radius'] = radius
  155. node_parent['x_var'] = x_var
  156. node_parent['y_var'] = y_var
  157.  
  158. node = node_parent
  159. node_parent = node['parent']
  160. else:
  161. #len(children) == max_elements_per_cluster => The leaf must be split
  162.  
  163. #Splits along the direction with highest variance
  164. if node['x_var'] >= node['y_var']:
  165. points.sort(key=lambda p: p['x'])
  166. else:
  167. points.sort(key=lambda p: p['y'])
  168.  
  169. #The new nodes have exactly half the elements of the old one
  170. new_node_1 = {'points': points[:split_size], 'leaf': True}
  171. new_node_2 = {'points': points[split_size:], 'leaf': True}
  172.  
  173.  
  174. #Compute the centroids for the new nodes
  175. for new_node in [new_node_1, new_node_2]:
  176. points = new_node['points']
  177. x_node = 0.
  178. y_node = 0.
  179. for point in points:
  180. x_node += point['x']
  181. y_node += point['y']
  182. n_p = len(points)
  183. x_node /= n_p
  184. y_node /= n_p
  185.  
  186. new_node['x'] = x_node
  187. new_node['y'] = y_node
  188.  
  189. #Adds the new point to the one of the two new nodes that is closest to the old centroid
  190. x_node = node['x']
  191. y_node = node['y']
  192. dist_1 = (x_node - new_node_1['x']) ** 2 + (y_node - new_node_1['y']) ** 2
  193. dist_2 = (x_node - new_node_2['x']) ** 2 + (y_node - new_node_2['y']) ** 2
  194.  
  195. if (dist_1 > dist_2):
  196. new_node = new_node_2
  197. new_node_2 = new_node_1
  198. new_node_1 = new_node
  199.  
  200. #INVARIANT: at this point new_node_1 is the one of the two new nodes closest to the old node's centroid
  201. #Adds the new point to new_node_1
  202. points = new_node_1['points']
  203. n_p = len(points)
  204. #Updates new_node_1's centroid
  205. x_node = new_node_1['x']
  206. y_node = new_node_1['y']
  207. x_node *= n_p
  208. y_node *= n_p
  209. x_node += new_point['x']
  210. y_node += new_point['y']
  211. points.append(new_point)
  212. n_p += 1
  213. new_node_1['x'] = x_node / n_p
  214. new_node_1['y'] = y_node / n_p
  215.  
  216. #Compute the radius of the new nodes
  217. for new_node in [new_node_1, new_node_2]:
  218.  
  219. x_node = new_node['x']
  220. y_node = new_node['y']
  221.  
  222. radius = 0.
  223. x_var = y_var = 0.
  224. for point in new_node['points']:
  225. #INVARIANT: point don't have radius
  226. x_dist = (x_node - point['x']) ** 2
  227. y_dist = (y_node - point['y']) ** 2
  228. radius = max(radius, x_dist + y_dist)
  229. #We don't need the exact variance, we can do fine with an estimate based on max distance form the centroid
  230. x_var = max(x_var, x_dist)
  231. y_var = max(y_var, y_dist)
  232.  
  233. new_node['radius'] = sqrt(radius)
  234. new_node['x_var'] = x_var
  235. new_node['y_var'] = y_var
  236.  
  237.  
  238. #INVARIANT: at this new_point new_node_1 is the closest to the centroid of node, so it takes its place among the
  239. #childrens of its parent
  240. node_parent = node['parent']
  241.  
  242. if node_parent == None:
  243. #The node that has just been split was the root: so it must create a new root...
  244. tree = {'children': [new_node_1, new_node_2], 'leaf':False, 'parent': None,
  245. 'x': (new_node_1['x'] + new_node_2['x'])/2,
  246. 'y': (new_node_1['y'] + new_node_2['y'])/2}
  247. x_dist_1 = (new_node_1['x'] - tree['x']) ** 2
  248. x_dist_2 = (new_node_2['x'] - tree['x']) ** 2
  249. y_dist_1 = (new_node_1['y'] - tree['y']) ** 2
  250. y_dist_2 = (new_node_2['y'] - tree['y']) ** 2
  251. tree['radius'] = max(sqrt(x_dist_1 + y_dist_1) + new_node_1['radius'],
  252. sqrt(x_dist_2 + y_dist_2) + new_node_2['radius'])
  253. tree['x_var'] = max(x_dist_1 + new_node_1['radius'] ** 2,
  254. x_dist_2 + new_node_2['radius'] ** 2)
  255. tree['y_var'] = max(y_dist_1 + new_node_1['radius'] ** 2,
  256. y_dist_2 + new_node_2['radius'] ** 2)
  257.  
  258. new_node_1['parent'] = new_node_2['parent'] = tree
  259.  
  260. #... and return it
  261. return tree
  262. else:
  263. #Replaces the old node (the one just split) with the closest of the newly created
  264. new_node_1['parent'] = node_parent
  265.  
  266. node_parent['children'].remove(node)
  267. node_parent['children'].append(new_node_1)
  268.  
  269.  
  270. while node_parent != None:
  271. node = node_parent
  272. children = node['children']
  273.  
  274. #Checks if there is still a node resulting from the split of one of its children
  275. #INVARIANT: new_node_2 is the farthest of the two resulting node from the split
  276. if new_node_2:
  277.  
  278. if len(children) < max_elements_per_cluster:
  279. #No need for farther splits: just append the new node
  280. children.append(new_node_2)
  281. new_node_2['parent'] = node
  282. new_node_2 = None
  283. else:
  284. #Must split this node too
  285. old_node = new_node_2
  286.  
  287. #Split the children along the axes with the biggest variance
  288. if node['x_var'] >= node['y_var']:
  289. children.sort(key=lambda p: p['x'])
  290. else:
  291. children.sort(key=lambda p: p['y'])
  292.  
  293. new_children = children[:split_size]
  294. new_node_1 = {'children': new_children, 'leaf': node['leaf']}
  295. for child in new_children:
  296. child['parent'] = new_node_1
  297.  
  298. new_children = children[split_size:]
  299. new_node_2 = {'children': new_children, 'leaf': node['leaf']}
  300. for child in new_children:
  301. child['parent'] = new_node_2
  302.  
  303. #Compute the centroids
  304. for new_node in [new_node_1, new_node_2]:
  305. x_node = 0.
  306. y_node = 0.
  307. for child in new_node['children']:
  308. x_node += child['x']
  309. y_node += child['y']
  310. n_p = len(new_node['children'])
  311. new_node['x'] = x_node / n_p
  312. new_node['y'] = y_node / n_p
  313.  
  314. #Finds the one of the new nodes closest to the original centroid
  315. dist_1 = (node['x'] - new_node_1['x']) ** 2 + (node['y'] - new_node_1['y']) ** 2
  316. dist_2 = (node['x'] - new_node_2['x']) ** 2 + (node['y'] - new_node_2['y']) ** 2
  317.  
  318. if (dist_1 > dist_2):
  319. new_node = new_node_2
  320. new_node_2 = new_node_1
  321. new_node_1 = new_node
  322.  
  323. #INVARIANT: At this point new_node_1 is the one of two nodes resulting from the split
  324. # closest to the orginal centroid
  325. n_p = len(new_node_1['children'])
  326. new_node_1['children'].append(old_node)
  327. old_node['parent'] = new_node_1
  328.  
  329. x_node = new_node_1['x']
  330. y_node = new_node_1['y']
  331. x_node *= n_p
  332. y_node *= n_p
  333. x_node += old_node['x']
  334. y_node += old_node['y']
  335. n_p += 1
  336. new_node_1['x'] = x_node / n_p
  337. new_node_1['y'] = y_node / n_p
  338.  
  339. #Compute the radiuses and the variances
  340. for new_node in [new_node_1, new_node_2]:
  341.  
  342. x_node = new_node['x']
  343. y_node = new_node['y']
  344.  
  345. radius = 0.
  346. x_var = y_var = 0.
  347.  
  348. for child_node in new_node['children']:
  349. x_dist = (x_node - child_node['x']) ** 2
  350. y_dist = (y_node - child_node['y']) ** 2
  351. radius = max(radius, sqrt(x_dist + y_dist) + child_node['radius'])
  352. #We don't need the exact variance, we can do fine with an estimate based on max distance form the centroid
  353. x_var = max(x_var, x_dist + child_node['radius'] ** 2)
  354. y_var = max(y_var, y_dist + child_node['radius'] ** 2)
  355.  
  356. new_node['radius'] = radius
  357. new_node['x_var'] = x_var
  358. new_node['y_var'] = y_var
  359.  
  360. #Checks whether the root has been split
  361. node_parent = node['parent']
  362. if node_parent == None:
  363. #Has just split the root
  364. tree = {'children': [new_node_1, new_node_2], 'leaf':False, 'parent': None,
  365. 'x': (new_node_1['x'] + new_node_2['x'])/2,
  366. 'y': (new_node_1['y'] + new_node_2['y'])/2}
  367. x_dist_1 = (new_node_1['x'] - tree['x']) ** 2
  368. x_dist_2 = (new_node_2['x'] - tree['x']) ** 2
  369. y_dist_1 = (new_node_1['y'] - tree['y']) ** 2
  370. y_dist_2 = (new_node_2['y'] - tree['y']) ** 2
  371. tree['radius'] = max(sqrt(x_dist_1 + y_dist_1) + new_node_1['radius'],
  372. sqrt(x_dist_2 + y_dist_2) + new_node_2['radius'])
  373. tree['x_var'] = max(x_dist_1 + new_node_1['radius'] ** 2, x_dist_2 + new_node_2['radius'] ** 2)
  374. tree['y_var'] = max(y_dist_1 + new_node_1['radius'] ** 2, y_dist_2 + new_node_2['radius'] ** 2)
  375. new_node_1['parent'] = new_node_2['parent'] = tree
  376. return tree
  377. else:
  378. new_node_1['parent'] = node_parent
  379.  
  380. node_parent['children'].remove(node)
  381. node_parent['children'].append(new_node_1)
  382.  
  383. #node doesn't exist anymore, and for new_node_1 and new_node_2 everything has been computed
  384. #and therefore can go to the next iteration
  385. continue
  386.  
  387. #Updates node's centroid, radius and variances
  388. x_node = 0.
  389. y_node = 0.
  390.  
  391. for child_node in children:
  392. x_node += child_node['x']
  393. y_node += child_node['y']
  394.  
  395. n_p = len(children)
  396. x_node /= n_p
  397. y_node /= n_p
  398. node['x'] = x_node
  399. node['y'] = y_node
  400.  
  401. radius = 0.
  402. x_var = y_var = 0.
  403. for child_node in children:
  404. x_dist = (x_node - child_node['x']) ** 2
  405. y_dist = (y_node - child_node['y']) ** 2
  406. radius = max(radius, sqrt(x_dist + y_dist) + child_node['radius'])
  407. x_var = max(x_var, x_dist + child_node['radius'] ** 2)
  408. y_var = max(y_var, y_dist + child_node['radius'] ** 2)
  409.  
  410. node['radius'] = radius
  411. node['x_var'] = x_var
  412. node['y_var'] = y_var
  413.  
  414. node_parent = node['parent']
  415.  
  416. return tree
  417.  
  418. #Creates the SS-tree for the topics (points)
  419. topics_tree = ss_make_tree()
  420.  
  421. #List of the questions for which a topic is relevant
  422. topics_relevant_questions = {}
  423.  
  424. #Reads the topics list
  425. regex = re.compile(TOPIC_REGEXP)
  426. for i in range(T):
  427. line = f.readline()
  428. m = regex.match(line)
  429. t_id = int(m.group(1))
  430. (x,y) = map(lambda s: float(s), m.group(2,3))
  431. topics[t_id] = (x, y)
  432. #Add the topic to the SS-tree
  433. topics_tree = ss_tree_insert({'x':x, 'y':y, 't_id':t_id}, topics_tree)
  434. #List of the questions for which a topic is relevant (initializes it)
  435. topics_relevant_questions[t_id] = []
  436.  
  437.  
  438. #Reads the questions list
  439. regex = re.compile(INTEGER_RE)
  440. for i in range(Q):
  441. line = f.readline()
  442. m = regex.findall(line)
  443.  
  444. q_id = int(m[0])
  445. Qn = int(m[1])
  446.  
  447. if (Qn!=0):
  448. Qids = map(lambda s: int(s), m[2:Qn+2]) #could have been [2:len(m)], but it is better to trigger an exception if the input is not well formed
  449. questions.append((q_id, Qids))
  450. for t_id in Qids:
  451. topics_relevant_questions[t_id].append(q_id)
  452.  
  453. #Uses a statically allocated array of doubles for the distances' heap, to avoid runtime checks and improve performance
  454. heap = array('d', range(100)) #INVARIANT: no more than 100 results are required for each query
  455.  
  456.  
  457. #Reads and processes the queries list
  458. regex = re.compile(QUERY_REGEXP)
  459. for i in range(N):
  460. line = f.readline()
  461. m = regex.match(line)
  462. q_type = m.group(1)
  463. n_res = int(m.group(2))
  464. if n_res == 0:
  465. print ''
  466. #DEBUG file_out.write('\n')
  467. continue
  468.  
  469. (x0,y0) = map(lambda s: float(s), m.group(3,4))
  470.  
  471. ''' Compares two elements of the solution: each element is a tuple (distance, ID),
  472. If the two distances are within a tolerance of 0.001 they are (by specs)
  473. considered equals, so in that case the couple with the highest ID is
  474. smaller; otherwise the order is determined by the two distances.
  475.  
  476. @param (da,ia): The first tuple to compare;
  477. @param (db,ib): The second tuple to compare;
  478. @return: An integer n:
  479. < 0 <=> (da,ia) < (db,ib) [da-db < -0.001 or ia > ib]
  480. > 0 <=> (da,ia) > (db,ib) [da-db > 0.001 or ia < ib]
  481. 0 <=> (da,ia) == (db,ib) [fabs(da-db) < 0.001 and ia == ib]
  482. '''
  483. def compare_items((da,ia), (db,ib)):
  484. #checks the distances first: must be greater than the threshold (distances are non-negatives!)
  485. if da < db - 0.001:
  486. return -1
  487. elif da > db + 0.001:
  488. return 1
  489. else:
  490. #if the distances are within threshold, then compares ids
  491. return ib - ia
  492.  
  493. #Switches the type of query
  494. if q_type == 't':
  495. #Init the heap to an empty max-heap
  496. heap_size = 0
  497. #Keeps track of the candidates to nearest neighbours found
  498. heap_elements = []
  499.  
  500. #Starts a search in the topics SS-tree;
  501. #All the topics are pushed in a bounded max-heap which holds at most n_res distances
  502. #(the n_res smallest ones) so that, once the heap is full, its first element is
  503. #the kth distance discovered so var, and this value can be used to prune the search
  504. #on the SS-tree.
  505.  
  506. if topics_tree['leaf']:
  507. #The tree has only one node, the root: so every point must be examined
  508. points = topics_tree['points']
  509. for p in points:
  510. t_id = p['t_id']
  511. x = p['x']
  512. y = p['y']
  513.  
  514. new_dist = sqrt((x - x0) ** 2 + (y - y0) ** 2)
  515.  
  516. if heap_size == n_res:
  517. if new_dist > heap[0]:
  518. #The heap is full: if the new value is greather than the kth distance,
  519. #then it can't be one of the k nearest neighbour's distances
  520. continue
  521.  
  522. heap_elements.append((new_dist, t_id))
  523. pos = 0
  524. # Bubble up the greater child until hitting a leaf.
  525. child_pos = 2 * pos + 1 # leftmost child position
  526. while child_pos < heap_size:
  527. # Set childpos to index of greater child.
  528. right_pos = child_pos + 1
  529. if right_pos < heap_size and heap[child_pos] < heap[right_pos]:
  530. child_pos = right_pos
  531. # Move the greater child up.
  532. if heap[child_pos] <= new_dist:
  533. break
  534. heap[pos] = heap[child_pos]
  535. pos = child_pos
  536. child_pos = 2*pos + 1
  537. heap[pos] = new_dist
  538. else:
  539. heap_elements.append((new_dist, t_id))
  540. heap[heap_size] = new_dist
  541. pos = heap_size
  542. heap_size += 1
  543. # Follow the path to the root, moving parents down until finding a place
  544. # newitem fits.
  545. while pos > 0:
  546. parent_pos = (pos - 1) >> 1
  547. parent = heap[parent_pos]
  548. if new_dist > parent:
  549. heap[pos] = parent
  550. pos = parent_pos
  551. else:
  552. break
  553. heap[pos] = new_dist
  554. else:
  555. queue = []
  556. #Adds all the root's children to the queue, and examines them in order of increasing distance
  557. #of their border from the query point
  558. children = topics_tree['children']
  559. for child in children:
  560. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  561. radius = child['radius']
  562. if dist <= radius:
  563. dist = 0
  564. else:
  565. dist -= radius
  566. queue.append((dist, radius, child))
  567.  
  568. queue.sort(key=lambda q:q[0], reverse=True)
  569.  
  570. while len(queue) > 0:
  571. (d, r, node) = queue.pop()
  572.  
  573. if node['leaf']:
  574. points = node['points']
  575. for p in points:
  576. t_id = p['t_id']
  577. x = p['x']
  578. y = p['y']
  579.  
  580. new_dist = sqrt((x - x0) ** 2 + (y - y0) ** 2)
  581.  
  582. if heap_size == n_res:
  583. #The heap is full: if the new value is greather than the kth distance,
  584. #then it can't be one of the k nearest neighbour's distances
  585. if new_dist > heap[0]:
  586. continue
  587.  
  588. heap_elements.append((new_dist, t_id))
  589. #heap[0] = new_dist
  590. pos = 0
  591. # Bubble up the greater child until hitting a leaf.
  592. child_pos = 2 * pos + 1 # leftmost child position
  593. while child_pos < heap_size:
  594. # Set childpos to index of greater child.
  595. right_pos = child_pos + 1
  596. if right_pos < heap_size and heap[child_pos] < heap[right_pos]:
  597. child_pos = right_pos
  598. # Move the greater child up.
  599. if heap[child_pos] <= new_dist:
  600. break
  601. heap[pos] = heap[child_pos]
  602. pos = child_pos
  603. child_pos = 2*pos + 1
  604. heap[pos] = new_dist
  605. else:
  606. heap_elements.append((new_dist, t_id))
  607. heap[heap_size] = new_dist
  608. pos = heap_size
  609. heap_size += 1
  610. # Follow the path to the root, moving parents down until it finds a place
  611. #where new_item fits.
  612. while pos > 0:
  613. parent_pos = (pos - 1) >> 1
  614. parent = heap[parent_pos]
  615. if new_dist > parent:
  616. heap[pos] = parent
  617. pos = parent_pos
  618. else:
  619. break
  620. heap[pos] = new_dist
  621.  
  622. #Checks if now the queue is full
  623. if heap_size == n_res:
  624. #If it is so, filters the queue
  625. #The heap is full: if the distance of the border of the node from the query point
  626. #is greather than the kth distance then no point in that node can be one of the
  627. #k nearest neighbour's
  628. d_max = heap[0]
  629. queue = [(d, r, n) for (d, r, n) in queue if d <= d_max]
  630. else:
  631. if heap_size < n_res:
  632. for child in node['children']:
  633. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  634. radius = child['radius']
  635. if dist <= radius:
  636. dist = 0
  637. else:
  638. dist -= radius
  639. queue.append((dist, radius, child))
  640.  
  641. queue.sort(key=lambda q:q[0], reverse=True)
  642. else:
  643. d_max = heap[0]
  644. queue = [(d, r, n) for (d, r, n) in queue if d <= d_max]
  645. for child in node['children']:
  646. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  647. radius = child['radius']
  648. if dist <= radius:
  649. dist = 0
  650. else:
  651. dist -= radius
  652.  
  653. if dist <= d_max:
  654. #The heap is full: if the distance of the border of the node from the query point
  655. #is greather than the kth distance then no point in that node can be one of the
  656. #k nearest neighbour's
  657. queue.append((dist, radius, child))
  658.  
  659. queue = sorted([(d, r, n) for (d, r, n) in queue if d <= d_max],
  660. key=lambda q:q[0], reverse=True)
  661.  
  662. #Filters the possible nearest neighbours such that their distance is not greater than the the distance of the kth
  663. #nearest neighbour (plus the tolerance)
  664. s = ['{} '.format(i_d) for (d, i_d) in
  665. sorted([(d, i_d) for (d, i_d) in heap_elements if d <= heap[0] + 0.001],
  666. cmp=compare_items)[:n_res]]
  667.  
  668. print ''.join(s)
  669. #DEBUG
  670. # file_out.write(''.join(s))
  671. # file_out.write('\n')
  672. else:
  673. #query type 'q'
  674.  
  675. #Starts a query on the topics, and as soon as it encounters new topics
  676. #(from the closest ones to the query point to the farthest)
  677.  
  678. questions_heap = []
  679. heap_elements = {}
  680.  
  681. if topics_tree['leaf']:
  682. #The SS-tree for topics is just made of a root node:
  683. #All the points must be checked
  684.  
  685. points = topics_tree['points']
  686. for p in points:
  687. t_id = p['t_id']
  688. (x,y) = topics[t_id]
  689.  
  690. new_dist = sqrt((x - x0) ** 2 + (y - y0) ** 2)
  691.  
  692. if len(heap_elements) >= n_res:
  693. if new_dist > questions_heap[n_res-1] + 0.001:
  694. continue
  695.  
  696. useful = False
  697. for q_id in topics_relevant_questions[t_id]:
  698. if q_id in heap_elements and heap_elements[q_id] <= new_dist:
  699. continue
  700. else:
  701. useful = True
  702. heap_elements[q_id] = new_dist
  703.  
  704. if useful:
  705. questions_heap = sorted(heap_elements.itervalues())
  706. else:
  707.  
  708. for q_id in topics_relevant_questions[t_id]:
  709. if q_id in heap_elements and heap_elements[q_id] <= new_dist:
  710. continue
  711. else:
  712. heap_elements[q_id] = new_dist
  713. if len(heap_elements) >= n_res:
  714. questions_heap = sorted(heap_elements.itervalues())
  715. else:
  716. queue = []
  717.  
  718. #Adds all the root's children to the queue, and then examine them one by one from the closest points to the
  719. #query point to the farthest ones: for each one checks the queries for which the topic is relevant
  720. #and for each of this queries checks if other closer relevant topics have been already found or not.
  721. #When at list n_res different questions have been met, starts comparing the distance from the query point
  722. #of farthest one to the distances (from the query point) of the SS-tree nodes' borders, pruning the search
  723. #on the nodes too far away.
  724. children = topics_tree['children']
  725. for child in children:
  726. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  727. radius = child['radius']
  728. if dist <= radius:
  729. dist = 0
  730. else:
  731. dist -= radius
  732. queue.append((dist, radius, child))
  733.  
  734. queue.sort(key=lambda q:q[0], reverse=True)
  735.  
  736. while len(queue) > 0:
  737. (d, r, node) = queue.pop()
  738.  
  739. if node['leaf']:
  740. points = node['points']
  741. for p in points:
  742. t_id = p['t_id']
  743. (x,y) = topics[t_id]
  744.  
  745. new_dist = sqrt((x - x0) ** 2 + (y - y0) ** 2)
  746.  
  747. if len(heap_elements) >= n_res:
  748. if new_dist > questions_heap[n_res-1] + 0.001:
  749. continue
  750.  
  751. useful = False
  752. for q_id in topics_relevant_questions[t_id]:
  753. if q_id in heap_elements and heap_elements[q_id] <= new_dist:
  754. continue
  755. else:
  756. useful = True
  757. heap_elements[q_id] = new_dist
  758.  
  759. if useful:
  760. questions_heap = sorted(heap_elements.itervalues())
  761. else:
  762.  
  763. for q_id in topics_relevant_questions[t_id]:
  764. if q_id in heap_elements and heap_elements[q_id] <= new_dist:
  765. continue
  766. else:
  767. heap_elements[q_id] = new_dist
  768.  
  769. if len(heap_elements) >= n_res:
  770. questions_heap = sorted(heap_elements.itervalues())
  771.  
  772. #Checks if it has already found at least n_res questions
  773. if len(heap_elements) >= n_res:
  774. #If it is so, filters the queue
  775. d_max = questions_heap[n_res-1] + 0.001
  776. queue = [(d, r, n) for (d, r, n) in queue if d <= d_max]
  777. else:
  778.  
  779. if len(heap_elements) < n_res:
  780. for child in node['children']:
  781. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  782. radius = child['radius']
  783. if dist <= radius:
  784. dist = 0
  785. else:
  786. dist -= radius
  787. queue.append((dist, radius, child))
  788.  
  789. queue.sort(key=lambda q:q[0], reverse=True)
  790. else:
  791. for child in node['children']:
  792. dist = sqrt((child['x'] - x0) ** 2 + (child['y'] - y0) ** 2)
  793. radius = child['radius']
  794. if dist <= radius:
  795. dist = 0
  796. else:
  797. dist -= radius
  798. d_max = questions_heap[n_res-1] + 0.001
  799. if dist <= d_max:
  800. queue.append((dist, radius, child))
  801.  
  802. queue = sorted([(d, r, n) for (d, r, n) in queue if d <= d_max], key=lambda q:q[0], reverse=True)
  803.  
  804.  
  805.  
  806.  
  807. s = ['{} '.format(i_d) for (d, i_d) in
  808. sorted([(d, i_d) for (i_d, d) in heap_elements.iteritems()],
  809. cmp=compare_items)[:n_res] ]
  810.  
  811. print ''.join(s)
  812. #DEBUG
  813. # file_out.write(''.join(s))
  814. # file_out.write('\n')
  815.  
  816. return
  817.  
  818. if __name__ == '__main__':
  819.  
  820. read_and_process_input(stdin)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement