fork download
copy 
  1. import bisect
  2.  
  3. def main():
  4.     import sys
  5.     input = sys.stdin.read
  6.     data = input().split()
  7.     ptr = 0
  8.     N = int(data[ptr])
  9.     ptr += 1
  10.     K = int(data[ptr])
  11.     ptr += 1
  12.     Q = int(data[ptr])
  13.     ptr += 1
  14.  
  15.     A = [0] * N
  16.  
  17.     sorted_values = []
  18.     freq = {}
  19.  
  20.     def add_value(v):
  21.         if v not in freq:
  22.             bisect.insort(sorted_values, -v)  # Store in descending order using negatives
  23.             freq[v] = 0
  24.         freq[v] += 1
  25.  
  26.     def remove_value(v):
  27.         freq[v] -= 1
  28.         if freq[v] == 0:
  29.             idx = bisect.bisect_left(sorted_values, -v)
  30.             del sorted_values[idx]
  31.             del freq[v]
  32.  
  33.     # Initialize all zeros
  34.     add_value(0)
  35.     freq[0] += (N - 1)  # Initially, all zeros; add N-1 more (1 added already)
  36.  
  37.     for _ in range(Q):
  38.         X_i = int(data[ptr]) - 1  # Convert to 0-based
  39.         ptr += 1
  40.         Y_i = int(data[ptr])
  41.         ptr += 1
  42.  
  43.         old_val = A[X_i]
  44.         new_val = Y_i
  45.  
  46.         # Remove old_val
  47.         if freq[old_val] == 1:
  48.             idx = bisect.bisect_left(sorted_values, -old_val)
  49.             del sorted_values[idx]
  50.             del freq[old_val]
  51.         else:
  52.             freq[old_val] -= 1
  53.  
  54.         # Add new_val
  55.         if new_val in freq:
  56.             freq[new_val] += 1
  57.         else:
  58.             bisect.insort(sorted_values, -new_val)
  59.             freq[new_val] = 1
  60.  
  61.         A[X_i] = new_val
  62.  
  63.         # Compute sum of top K elements
  64.         current_sum = 0
  65.         count = 0
  66.         for val in sorted_values:
  67.             v = -val
  68.             cnt = freq[v]
  69.             if count + cnt <= K:
  70.                 current_sum += v * cnt
  71.                 count += cnt
  72.             else:
  73.                 current_sum += v * (K - count)
  74.                 count = K
  75.                 break
  76.             if count >= K:
  77.                 break
  78.         print(current_sum)
  79.  
  80. if __name__ == '__main__':
  81.     main()
Success #stdin #stdout 0.02s 7312KB
comments (?)
stdin
copy 
4 2 10
1 5
2 1
3 3
4 2
2 10
1 0
4 0
3 1
2 0
3 0
stdout
copy 
5
6
8
8
15
13
13
11
1
0
https://ideone.com/oGEnjB
language:
Python (cpython 2.7.16)
created:
1 day 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!