fork download
  1. def getMinConnectionCost(warehouseCapacity, additionalHubs):
  2. n = len(warehouseCapacity)
  3. q = len(additionalHubs)
  4. result = []
  5.  
  6. # Pre-calculate prefix sums and costs to central hub
  7. prefix_sum = [0] * (n + 1)
  8. for i in range(n):
  9. prefix_sum[i + 1] = prefix_sum[i] + warehouseCapacity[i]
  10.  
  11. # Pre-calculate total cost to central hub
  12. total_to_n = 0
  13. for i in range(n-1):
  14. total_to_n += warehouseCapacity[n-1] - warehouseCapacity[i]
  15.  
  16. for hubs in additionalHubs:
  17. hubA, hubB = hubs[0] - 1, hubs[1] - 1 # Convert to 0-based indexing
  18.  
  19. # Calculate costs using prefix sums
  20. total_cost = 0
  21.  
  22. # Segment 1: Before hubA
  23. if hubA > 0:
  24. count = hubA
  25. total_cost += count * warehouseCapacity[hubA] - (prefix_sum[hubA] - prefix_sum[0])
  26.  
  27. # Segment 2: Between hubA and hubB
  28. if hubB > hubA + 1:
  29. count = hubB - (hubA + 1)
  30. total_cost += count * warehouseCapacity[hubB] - (prefix_sum[hubB] - prefix_sum[hubA + 1])
  31.  
  32. # Segment 3: After hubB to before n
  33. if hubB < n-2:
  34. total_cost += total_to_n - (prefix_sum[n-1] - prefix_sum[hubB + 1]) - \
  35. (n-2 - hubB) * warehouseCapacity[n-1]
  36.  
  37. result.append(total_cost)
  38.  
  39. return result
  40.  
  41. # Test cases
  42. def test():
  43. # Sample Case 0
  44. warehouseCapacity0 = [0, 2, 5, 9, 12, 18]
  45. additionalHubs0 = [[2, 5], [1, 3]]
  46. print(getMinConnectionCost(warehouseCapacity0, additionalHubs0)) # Expected: [12, 18]
  47.  
  48. # Sample Case 1
  49. warehouseCapacity1 = [2, 6, 8, 14]
  50. additionalHubs1 = [[1, 2]]
  51. print(getMinConnectionCost(warehouseCapacity1, additionalHubs1)) # Expected: [6]
  52.  
  53. if __name__ == "__main__":
  54. test()
Success #stdin #stdout 0.09s 14152KB
stdin
Standard input is empty
stdout
[12, 8]
[4]