import sys, math, random
def read_int_list():
return list(map(int, [x for x in sys.stdin.readline().strip().split(' ') if x != '']))
def computeResult(A):
Z = [0]*n
for i in range(n):
if A[i] > 0:
Z[i] = 1
while True:
# print(Z)
X = [Z[0]]
for i in range(1, n):
X.append(X[-1] + Z[i])
max_b = 0
i_max = None
i_relax = None
for i in range(n):
x = X[i-1] if i > 0 else 0
z = Z[i]
y = X[n-1] - X[i]
b = A[i] - x*y*z - (x+y)*z*(z-1)//2 - z*(z-1)*(z-2)//6
if b/(1+Z[i]) > max_b:
max_b = b
i_max = i
if i_max is None:
break
Z[i_max] += 1
while True:
relax = False
for ix in range(n):
if not((A[ix] > 0 and Z[ix] > 2) or (A[ix] == 0 and Z[ix] > 1)):
continue
Z[ix] -= 1
X = [Z[0]]
for i in range(1, n):
X.append(X[-1] + Z[i])
B = [0]*n
for i in range(n):
x = X[i-1] if i > 0 else 0
z = Z[i]
y = X[n-1] - X[i]
B[i] = x*y*z + (x+y)*z*(z-1)//2 + z*(z-1)*(z-2)//6
relax = True
for i in range(n):
if B[i] < A[i]:
relax = False
break
if relax:
break
else:
Z[ix] += 1
if not relax:
break
res = sum(Z)
# print(Z)
return res
def verify():
n = 5
Z = [0]*n
for i in range(n):
Z[i]=random.randint(0,10)
#Z = [3, 7, 1, 10, 9]
Z = [2, 10, 10, 9, 10]
X = [Z[0]]
for i in range(1, n):
X.append(X[-1] + Z[i])
B = [0]*n
for i in range(n):
x = X[i-1] if i > 0 else 0
z = Z[i]
y = X[n-1] - X[i]
B[i] = x*y*z + (x+y)*z*(z-1)//2 + z*(z-1)*(z-2)//6
print(Z)
print(B)
res = computeResult(B)
print(res, sum(Z))
if __name__=='__main__':
nTC = int(sys.stdin.readline().strip())
for tc in range(nTC):
n = int(sys.stdin.readline().strip())
A = read_int_list()
res = computeResult(A)
print(res)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import sys, math, random def read_int_list(): return list(map(int, [x for x in sys.stdin.readline().strip().split(' ') if x != ''])) def computeResult(A): Z = [0]*n for i in range(n): if A[i] > 0: Z[i] = 1 while True: # print(Z) X = [Z[0]] for i in range(1, n): X.append(X[-1] + Z[i]) max_b = 0 i_max = None i_relax = None for i in range(n): x = X[i-1] if i > 0 else 0 z = Z[i] y = X[n-1] - X[i] b = A[i] - x*y*z - (x+y)*z*(z-1)//2 - z*(z-1)*(z-2)//6 if b/(1+Z[i]) > max_b: max_b = b i_max = i if i_max is None: break Z[i_max] += 1 while True: relax = False for ix in range(n): if not((A[ix] > 0 and Z[ix] > 2) or (A[ix] == 0 and Z[ix] > 1)): continue Z[ix] -= 1 X = [Z[0]] for i in range(1, n): X.append(X[-1] + Z[i]) B = [0]*n for i in range(n): x = X[i-1] if i > 0 else 0 z = Z[i] y = X[n-1] - X[i] B[i] = x*y*z + (x+y)*z*(z-1)//2 + z*(z-1)*(z-2)//6 relax = True for i in range(n): if B[i] < A[i]: relax = False break if relax: break else: Z[ix] += 1 if not relax: break res = sum(Z) # print(Z) return res def verify(): n = 5 Z = [0]*n for i in range(n): Z[i]=random.randint(0,10) #Z = [3, 7, 1, 10, 9] Z = [2, 10, 10, 9, 10] X = [Z[0]] for i in range(1, n): X.append(X[-1] + Z[i]) B = [0]*n for i in range(n): x = X[i-1] if i > 0 else 0 z = Z[i] y = X[n-1] - X[i] B[i] = x*y*z + (x+y)*z*(z-1)//2 + z*(z-1)*(z-2)//6 print(Z) print(B) res = computeResult(B) print(res, sum(Z)) if __name__=='__main__': nTC = int(sys.stdin.readline().strip()) for tc in range(nTC): n = int(sys.stdin.readline().strip()) A = read_int_list() res = computeResult(A) print(res) |
English