fork download
  1. /* package whatever; // don't place package name! */
  2.  
  3. import java.util.*;
  4. import java.lang.*;
  5. import java.io.*;
  6.  
  7. /* Name of the class has to be "Main" only if the class is public. */
  8. class Ideone
  9. {
  10. public static void main (String[] args) throws java.lang.Exception
  11. {
  12. // your code goes here
  13. Scanner sc = new Scanner(System.in);
  14.  
  15. int m = sc.nextInt();
  16. int n = sc.nextInt();
  17.  
  18. int arr[] = new int[n+1];
  19. for(int i=1;i<=n;i++) arr[i]=sc.nextInt();
  20.  
  21. int arsi[] = new int[n+15];
  22. int krypto[] = new int[n+15];
  23.  
  24. int ac=0,kc=0;
  25.  
  26. for(int i=1;i<=n;i++){
  27. if(arr[i]==1 || arr[i]==3){
  28. arsi[i]=i;
  29. ac++;
  30. }
  31. if(arr[i]==2 || arr[i]==3){
  32. krypto[i]=i;
  33. kc++;
  34. }
  35. }
  36.  
  37. int prefixA[] = new int[n+15];
  38. int prefixK[] = new int[n+15];
  39.  
  40. for(int i=1;i<=n;i++){
  41. prefixA[i]=prefixA[i-1]+arsi[i];
  42. prefixK[i]=prefixK[i-1]+krypto[i];
  43. }
  44.  
  45. int suffixA[] = new int[n+15];
  46. int suffixK[] = new int[n+15];
  47.  
  48. for(int i=n;i>=1;i--){
  49. suffixA[i]=suffixA[i+1]+arsi[i];
  50. suffixK[i]=suffixK[i+1]+krypto[i];
  51. }
  52.  
  53. List<Integer> result = new ArrayList<>();
  54.  
  55. int ca=0,ck=0;
  56.  
  57. for(int i=1;i<=n;i++){
  58.  
  59. if(arsi[i]>0){
  60. ca++;
  61. }
  62.  
  63. int la = i*ca - prefixA[i];
  64. int lb = suffixA[i] - i*ac;
  65. int ra = la+lb;
  66.  
  67. if(arsi[i]>0) ac--;
  68.  
  69.  
  70. if(krypto[i]>0){
  71. ck++;
  72. }
  73.  
  74. int lk = i*ck - prefixK[i];
  75. int lka = suffixK[i] - i*kc;
  76. int rk = lk+lka;
  77.  
  78. if(krypto[i]>0) kc--;
  79.  
  80. result.add(Math.abs(rk-ra));
  81. }
  82.  
  83. System.out.println(result);
  84. }
  85. }
Success #stdin #stdout 0.12s 54472KB
stdin
1 5
2 1 3 0 1
stdout
[5, 2, 1, 0, 1]