fork(4) download 
  1. import time
  2. import random
  3.  
  4.  
  5. def distribute(arr, x):
  6.     total = sum(arr) + x
  7.     I = sorted(range(len(arr)), key=arr.__getitem__)
  8.     while I:
  9.         minimum, additional = divmod(total, len(I))
  10.         if arr[I[-1]] <= minimum:
  11.             break
  12.         total -= arr[I.pop()]
  13.     for i in sorted(I):
  14.         arr[i] = minimum
  15.         if additional > 0:
  16.             arr[i] += 1
  17.             additional -= 1
  18.  
  19. def f(A, x):
  20.   smallest = min(A)
  21.  
  22.   lo = smallest
  23.   hi = smallest + x
  24.  
  25.   while lo < hi:
  26.     mid = lo + (hi - lo) // 2
  27.  
  28.     can_reach = True
  29.     temp = x
  30.  
  31.     for a in A:
  32.       if a <= mid:
  33.         diff = mid - a
  34.         if diff > temp:
  35.           can_reach = False
  36.           break
  37.         else:
  38.           temp -= diff
  39.  
  40.     if can_reach:
  41.       lo = mid + 1
  42.     else:
  43.       hi = mid
  44.  
  45.   target = lo - 1
  46.  
  47.   for i, a in enumerate(A):
  48.     if a < target:
  49.       x -= target - a
  50.       A[i] = target
  51.  
  52.   if x:
  53.     for i, a in enumerate(A):
  54.       if a == target:
  55.         A[i] += 1
  56.         x -= 1
  57.         if x == 0:
  58.           break
  59.  
  60.   return A
  61.  
  62.  
  63. n = 1000000
  64.  
  65.  
  66. m = 20
  67. x = int(random.random() * m)
  68. A = [random.randint(0, m) for _ in range(n)]
  69. A1 = A.copy()
  70.  
  71. t0 = time.time()
  72. distribute(A, x)
  73. print(time.time() - t0)
  74.  
  75.  
  76. t0 = time.time()
  77. f(A1, x)
  78. print(time.time() - t0)
  79.  
  80. print(sorted(A) == sorted(A1))
  81.  
  82.  
  83. m = 100000000
  84. x = int(random.random() * m)
  85. A = [random.randint(0, m) for _ in range(n)]
  86. A1 = A.copy()
  87.  
  88. t0 = time.time()
  89. distribute(A, x)
  90. print(time.time() - t0)
  91.  
  92.  
  93. t0 = time.time()
  94. f(A1, x)
  95. print(time.time() - t0)
  96.  
  97. print(sorted(A) == sorted(A1))
Time limit exceeded #stdin #stdout 5s 113880KB
comments (?)
stdin 
Standard input is empty
stdout 
0.5914058685302734
0.07783126831054688
True
1.0516798496246338
0.5425384044647217
True
https://ideone.com/X6GxFD
language:
Python 3 (python  3.12)
created:
4 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!