I’ve beern trying to solve this question but getting timeout for most test cases. Can anyone help me in optimising this?
Problem Statement :
You are given an array A of length N. You have to choose a subset S from given array A, such that average of S is less than K. You need to print the maximum possible length of S.
Input format :
The first line of each input contains N, length of array A. Next line contains N space separated elements of array A. Next line of input contains an integer Q, Number of queries. Each following Q lines contains a single integer K.
Output format :
For each query, print single integer denoting the maximum possible length of the subset.
Sample Input
5 1 2 3 4 5 5 1 2 3 4 5
Sample Output
0 2 4 5 5
Explanation
- In first query, there is no possible subset such that its average is less than 1.
- In second query, you can select the subset {1,2}.
- In third query, you can select the subset {1,2,3,4}.
- In fourth and fifth query, you can seelct the complete array {1,2,3,4,5}.
Here’s my solution:
import java.util.*; public class Playground { public static void main(String args[] ) throws Exception { Scanner s = new Scanner(System.in); long n = Long.parseLong(s.nextLine()); // Reading input from STDIN String[] temp = s.nextLine().trim().split(" "); long[] arr = new long[(int) n]; for (int i = 0; i < n; i++) arr[i] = Integer.parseInt(temp[i]); long q = Long.parseLong(s.nextLine()); long[] queries = new long[(int) q]; for (int i = 0; i < q; i++) { long x = Long.parseLong(s.nextLine()); queries[i] = x; } PriorityQueue<Long> queue = new PriorityQueue<>(); for (long x : arr) queue.add(x); for (long x : queries) { double avg = 0; List<Long> list = new ArrayList<>(); int i = 0; int sum = 0; boolean flag = false; while (! queue.isEmpty()) { long num = queue.poll(); i++; list.add(num); sum += num; avg = (double) sum / i; if (avg >= x) { System.out.println(i - 1); flag = true; break; } } if (! flag) System.out.println(n); queue.addAll(list); } } }
Answer
An easy way to solve this is to sort the array first.
After you sorted the array so each element is equal or greater than the last, then solving a single run is easy:
int count = 0; int limit = 0; for (int i : sortedArray) { int diff = i - maxAvg; if (limit + diff < 0) { limit += diff; count++ } else { break; } } System.out.println(count);
This works because if the difference to the max average is negative you can use values with a positive difference until you hit the limit.
Sorting the array is O(n*log(n))
, and for each solution you only need O(n)
This is my full solution with all the parsing:
public static void main(String[] args) { Scanner sc = new Scanner(System.in); int arrLen = Integer.parseInt(sc.nextLine()); int[] array = new int[arrLen]; String[] strNums = sc.nextLine().split(" ", arrLen); for (int i = 0; i < arrLen; i++) { array[i] = Integer.parseInt(strNums[i]); } Arrays.sort(array); int numTests = Integer.parseInt(sc.nextLine()); for (int i = 0; i < numTests; i++) { int maxAvg = Integer.parseInt(sc.nextLine()); int limit = 0; int count = 0; for (int j : array) { int diff = j - maxAvg; if (limit + diff < 0) { count++; limit += diff; } else { break; } } System.out.println(count); } sc.close(); }