If anyone who was able to solve the problem tell me where my approach could've been optimized.

I took a map and put all scores which belong to university i in an array and stored it in the map against key i.

Then, I sorted the array for every key in the map, and transformed the array into prefix sum array.

My observation was that, say for university i, if the size of array against key i is not entirely divisible by k (k : 1 to n), then we have to ignore first k%n elements (n = size of array against key i in map), this translates to substraction of prefix sum at index k % n — 1 from entire sum of all students.

So basically, after this, I looped over all k (1 to n) and over all map keys, to find sizes which are indivisible by k and performed above operation.

This essentially gave me a TLE on test case 5.

I'd be really grateful for any tips/ help as to where I could further optimize this.

Thanks a lot.

HERE'S MY CODE:

public static void solve(int n, int[] u, long[] s){

Map<Integer,ArrayList<Long>> mp = new HashMap<>();

for(int i = 0 ; i < n ; i++){
  if(mp.containsKey(u[i]-1)){
    mp.get(u[i]-1).add(s[i]);
  }
  else{
    ArrayList<Long> ls = new ArrayList<>();
    ls.add(s[i]);
    mp.put(u[i]-1,ls);
  }
}

for(int i : mp.keySet()){
  Collections.sort(mp.get(i));
}

Long ans = 0l;
for(int i : mp.keySet()){
    for(long k : mp.get(i)){
      ans += k;
    }
}

for(int i : mp.keySet()){
  ArrayList<Long> ls = mp.get(i);
  for(int j = 1 ; j < ls.size() ; j++){
    ls.set(j,ls.get(j) + ls.get(j-1));
  }
}

Map<Integer,ArrayList<Integer>> sizes = new HashMap<>();

for(int k : mp.keySet()){
  if(sizes.containsKey(mp.get(k).size())){
    sizes.get(mp.get(k).size()).add(k);
  }
  else{
    ArrayList<Integer> ls = new ArrayList<>();
    ls.add(k);
    sizes.put(mp.get(k).size(),ls);
  }
}

System.out.print(ans + " ");
for(int i = 2 ; i <= n ; i++){
  long f = ans;
  for(int j : sizes.keySet()){
    if(j % i != 0){
      for(int k : sizes.get(j)){
          f -= mp.get(k).get((j % i)-1);
        }
      }
    }
  System.out.print(f + " ");
}
System.out.println("");

}