Submission #70460510

---

Source Code Expand

import sys
from collections import Counter
import numpy as np
# <nonumba> from numba import njit, jit, boolean, int64, float64, typeof

def main():
    from io import StringIO
    execute(StringIO(sample), sys.stderr)
    print("-- ^sample! --", file=sys.stderr)

    execute(sys.stdin, sys.stdout)

    # #<nonumba> if len(solve.nopython_signatures) != 1:
    # #<nonumba>     raise KeyboardInterrupt("!! Unexpected: wrong use of jit !!")

def execute(src, dst):
    N, M, C, = map(int, src.readline().split())
    A = list(map(int, src.readline().split()))
    # nda_1row = np.fromstring(src.readline().strip(), dtype=np.int64, sep=' ')
    cnter = Counter(A)
    BP = list(cnter.items())
    BP.sort()
    B = np.array([b for (b,p) in BP] , dtype=np.int64)
    P = np.array([p for (b,p) in BP] , dtype=np.int64)

    ans = solve(N,M,C,B,P)
    print(ans, file=dst)

sample = """\
5 3 2
1 2 1 0 1
"""

# #<nonumba> @jit(nopython=True, cache=True)
def solve(N,M,C,B,P):
    ans = 0
    G = len(B)
    prevB = -1
    seen = 0
    g = 1
    # stoppedB = 0
    stoppedH = 0
    for startB, startP in zip(B,P) :
        Bdiff = startB - prevB
        stop = False
        if startB == B[0]:
            seen = 0
        else:
            seen += - startP
        while True:
            # nextB = B[g]
            nextP = P[g]
            if C <= seen + nextP:
                while True:
                    if seen == C :
                        stop = True
                        rest = nextP - stoppedH
                        break
                    seen += 1
                    stoppedH += 1
                    continue
            if stop:
                break
            seen += nextP
            stoppedH = 0
            g = ( g + 1 ) % G
        Xi = C + nextP - stoppedH
        ans += Xi * Bdiff
        prevB = startB
        continue
    return ans




main()

Submission Info

Judge Result