fork download
  1. def merge_arrays(group1, group2):
  2. merged = []
  3. i = 0
  4. j = 0
  5.  
  6. while i < len(group1) and j < len(group2):
  7. if group1[i] <= group2[j]:
  8. merged.append((group1[i], 'group1'))
  9. i += 1
  10. else:
  11. merged.append((group2[j], 'group2'))
  12. j += 1
  13.  
  14. while i < len(group1):
  15. merged.append((group1[i], 'group1'))
  16. i += 1
  17.  
  18. while j < len(group2):
  19. merged.append((group2[j], 'group2'))
  20. j += 1
  21.  
  22. return merged
  23.  
  24. def calculate_component_sum(coords):
  25. total_sum = 0
  26. prefix_sum_group1 = 0
  27. prefix_sum_group2 = 0
  28. count_group1 = 0
  29. count_group2 = 0
  30.  
  31. for pos, group in coords:
  32. if group == 'group1':
  33. total_sum += count_group2 * pos - prefix_sum_group2
  34. prefix_sum_group1 += pos
  35. count_group1 += 1
  36. else:
  37. total_sum += count_group1 * pos - prefix_sum_group1
  38. prefix_sum_group2 += pos
  39. count_group2 += 1
  40.  
  41. return total_sum
  42.  
  43. def merge_and_calculate_safety(N, group1, Q, group2):
  44. merged_coords = merge_arrays(group1, group2)
  45. return calculate_component_sum(merged_coords)
Success #stdin #stdout 0.03s 9624KB
stdin
Standard input is empty
stdout
Standard output is empty