计算数组的最大长度，使平均值小于给定值

Question

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

1. In first query, there is no possible subset such that its average is less than 1.
2. In second query, you can select the subset {1,2}.
3. In third query, you can select the subset {1,2,3,4}.
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 1

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();
}