fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. int main() {
  5.     ios::sync_with_stdio(false);
  6.     cin.tie(nullptr);
  7.  
  8.     int n = 1, m;
  9.     cin >> m;
  10.  
  11.     vector<int> b(m+1), a(m+1,0), k(m+1,0);
  12.     for (int i=1; i<=m; i++) cin >> b[i];
  13.  
  14.     long long tot_k = 0, kc = 0;
  15.     for (int i = 1; i <= m; i++) {
  16.         if (b[i] == 1) {
  17.             a[i] = 1;
  18.         }
  19.         else if (b[i] == 2) {
  20.             k[i] = 1;
  21.             tot_k += i;
  22.             kc++;
  23.         }
  24.         else if (b[i] == 3) {
  25.             a[i] = 1;
  26.             k[i] = 1;
  27.             tot_k += i;
  28.             kc++;
  29.         }
  30.     }
  31.  
  32.     // ---- Aarsi optimized using prefix sums ----
  33.     vector<long long> prefixA(m+2,0), countA(m+2,0);
  34.     for (int i=1; i<=m; i++) {
  35.         prefixA[i] = prefixA[i-1] + (a[i] ? i : 0);
  36.         countA[i]  = countA[i-1] + (a[i] ? 1 : 0);
  37.     }
  38.     long long total_sumA = prefixA[m], total_countA = countA[m];
  39.  
  40.     vector<long long> aarsi(m+1,0);
  41.     for (int line=1; line<=m; line++) {
  42.         long long left_count = countA[line-1];
  43.         long long left_sum   = prefixA[line-1];
  44.         long long left_cost  = left_count * 1LL * line - left_sum;
  45.  
  46.         long long right_count = total_countA - countA[line];
  47.         long long right_sum   = total_sumA - prefixA[line];
  48.         long long right_cost  = right_sum - right_count * 1LL * line;
  49.  
  50.         aarsi[line] = left_cost + right_cost;
  51.     }
  52.  
  53.     // ---- Krypto costs ----
  54.     vector<long long> krypto(m+1,0);
  55.     long long sum_of_krypto = 0, count_of_krypto = 0;
  56.     for (int line=1; line<=m; line++) {
  57.         long long w = 0, d = 0;
  58.         if (k[line] == 1) {
  59.             sum_of_krypto += line;
  60.             w = line;
  61.             d = 1;
  62.             count_of_krypto++;
  63.         }
  64.         long long left  = count_of_krypto * 1LL * line - sum_of_krypto;
  65.         long long right = (tot_k - sum_of_krypto + w) - 1LL * line * (kc - count_of_krypto + d);
  66.         krypto[line] = left + right;
  67.     }
  68.  
  69.     // ---- Print results ----
  70.     for (int i=1; i<=m; i++) {
  71.         cout << llabs(aarsi[i] - krypto[i]) << " ";
  72.     }
  73.     cout << "\n";
  74.     return 0;
  75. }
  76.  
Success #stdin #stdout 0.01s 5276KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 

	

		
	 
	
	
			

					
			
			
			
			
			
			
			
			

			
			
			
		

	
		
		

			

			                        







			
			

				

	
					
					 https://ideone.com/MIgiQm
					
					
					
					
					

												

												
						
						

					

					
					
					

						
language:

						
C++ (gcc 8.3)

					

					

						
created:
								
						
11 days 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!