Advertisement
SqFKYo

Roller Coaster

Feb 7th, 2016
313
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.12 KB | None | 0 0
  1. import numpy as np
  2.  
  3. class FinishedException(Exception):
  4.     pass
  5.  
  6. spots, total_runs, groups_amount = [int(i) for i in input().split()]
  7.  
  8. groups = []
  9. for i in range(groups_amount):
  10.     groups.append(int(input()))
  11.    
  12. groups = np.array(groups)
  13.  
  14. # Initializing the money and runs
  15. money = np.array([0], dtype=np.int64)
  16. runs = total_runs
  17.  
  18. """
  19. # Very basic version, works if the dataset isn't huge.
  20. try:
  21.    while True:
  22.        cumulative = np.cumsum(groups)
  23.        found_index = np.searchsorted(cumulative, spots, side='right')
  24.        largest_found = cumulative[found_index-1]
  25.        money += largest_found
  26.        groups = np.roll(groups, -found_index)
  27.        runs -= 1
  28.        
  29.        if runs < 1:
  30.            raise FinishedException
  31. """
  32. # More advanced version, still doesn't work on large datasets
  33. # Initializing values
  34. cumulative = np.cumsum(groups)
  35. last_index = 0
  36. last_value = 0
  37.  
  38. try:
  39.     # If we can fit all the groups to the coaster at once, we get result from multiplication
  40.     total = cumulative[-1]
  41.     if total <= spots:
  42.         money += total*runs
  43.         raise FinishedException
  44.    
  45.     while True:
  46.         new_index = np.searchsorted(cumulative, last_value+spots, side='right')
  47.         new_value = cumulative[new_index-1]
  48.         # Updating cumulative to be the portion that hasn't been searched through yet
  49.         cumulative = cumulative[new_index:]
  50.        
  51.         new_size = cumulative.size
  52.        
  53.         if new_size == 0:
  54.             # We've reached end of the vector.
  55.             # Either we have right amount of people or not enough.
  56.             if (new_value - last_value + groups[0]) > spots:
  57.                 # Just the right amount of people, no need to roll,
  58.                 # But need to reinitialize the vectors
  59.                 money += (new_value - last_value)
  60.                 runs -= 1
  61.                 last_value = 0
  62.                 cumulative = np.cumsum(groups)
  63.                
  64.                 # At this point we've come around to the start
  65.                 # We check amount of runs done and left, and see if we can use
  66.                 # that info to remove tons of runs and add tons of money
  67.                 # with simple multiplication!
  68.                 divider = total_runs//(total_runs-runs)
  69.                 runs = total_runs-((total_runs-runs)*divider)
  70.                 money *= divider
  71.                
  72.             else:
  73.                 # Not quite enough people, making a vector by taking last_size
  74.                 # amount of components from end of 'groups' and then pasting 'groups'
  75.                 # in whole to it.
  76.                 last_value = 0
  77.                 cumulative = np.concatenate((groups[-last_size:], groups))
  78.                 cumulative = np.cumsum(cumulative)
  79.  
  80.         else:
  81.             # If not at the end of vector, adding the difference to the money amount
  82.             # and updating values
  83.             money += (new_value - last_value)
  84.             last_value = new_value
  85.             last_size = new_size
  86.             runs -= 1
  87.        
  88.         if runs < 1:
  89.             raise FinishedException
  90. except FinishedException:
  91.     print(money[0])
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement