Searching & SortingProblem 19 of 36
Bishu and Soldiers
Problem Statement
Bishu went to fight against N soldiers. Each soldier has a power denoted by an integer. Bishu has a power of X. He can kill all soldiers whose power is less than or equal to X. You are given Q queries, each containing a power value X. For each query, determine how many soldiers Bishu can kill and the total sum of their powers.
Constraints:
- 1 ≤ N ≤ 10^5
- 1 ≤ Q ≤ 10^5
- 1 ≤ Power of each soldier ≤ 10^6
- 1 ≤ X ≤ 10^6
Example:
- Input:
soldiers = [1, 2, 3, 4, 5, 6, 7],X = 3 - Output:
count = 3, sum = 6 - Explanation: Soldiers with power 1, 2, 3 can be killed. Sum = 1 + 2 + 3 = 6
Example 2:
- Input:
soldiers = [1, 2, 3, 4, 5, 6, 7],X = 10 - Output:
count = 7, sum = 28 - Explanation: All soldiers can be killed
Approach 1: Brute Force (Linear Search for Each Query)
Intuition
For each query, iterate through all soldiers and count those with power ≤ X while summing their powers.
Algorithm
- For each query X:
- Initialize count = 0, sum = 0
- Iterate through all soldiers
- If soldier power ≤ X, increment count and add to sum
- Return count and sum for each query
java
import java.util.*;
public class Solution {
public static long[] countAndSum(int[] soldiers, int X) {
int count = 0;
long sum = 0;
for (int power : soldiers) {
if (power <= X) {
count++;
sum += power;
}
}
return new long[]{count, sum};
}
public static void solveQueries(int[] soldiers, int[] queries) {
for (int X : queries) {
long[] result = countAndSum(soldiers, X);
System.out.println(result[0] + " " + result[1]);
}
}
}Complexity Analysis
- Time Complexity: O(q * n) - For each query, we scan all soldiers
- Space Complexity: O(1) - No extra space used
Approach 2: Sorting + Binary Search + Prefix Sum (Optimal)
Intuition
Sort the soldiers' powers and precompute prefix sums. For each query, use binary search to find the position up to which soldiers can be killed, then use prefix sum to get the total power.
Key Observations:
- Sorting allows us to use binary search to find the count
- Prefix sum gives us O(1) sum queries
- Binary search finds the rightmost position where power ≤ X
Algorithm
- Sort the soldiers array
- Compute prefix sum array
- For each query X:
- Use upper_bound to find count of soldiers with power ≤ X
- Use prefix sum to get total power
- Return results
java
import java.util.*;
public class BishuSoldiers {
private int[] soldiers;
private long[] prefixSum;
public BishuSoldiers(int[] arr) {
soldiers = arr.clone();
Arrays.sort(soldiers);
int n = soldiers.length;
prefixSum = new long[n + 1];
for (int i = 0; i < n; i++) {
prefixSum[i + 1] = prefixSum[i] + soldiers[i];
}
}
private int upperBound(int X) {
int left = 0, right = soldiers.length;
while (left < right) {
int mid = left + (right - left) / 2;
if (soldiers[mid] <= X) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
public long[] query(int X) {
int pos = upperBound(X);
return new long[]{pos, prefixSum[pos]};
}
public static void main(String[] args) {
Scanner sc = new Scanner(System.in);
int n = sc.nextInt();
int[] soldiers = new int[n];
for (int i = 0; i < n; i++) {
soldiers[i] = sc.nextInt();
}
BishuSoldiers bs = new BishuSoldiers(soldiers);
int q = sc.nextInt();
while (q-- > 0) {
int X = sc.nextInt();
long[] result = bs.query(X);
System.out.println(result[0] + " " + result[1]);
}
}
}Dry Run Example
soldiers = [7, 2, 5, 1, 6, 3, 4], X = 4
Step 1: Sort soldiers
sorted = [1, 2, 3, 4, 5, 6, 7]
Step 2: Compute prefix sum
prefixSum = [0, 1, 3, 6, 10, 15, 21, 28]
(0) (1) (1+2) (1+2+3) ...
Step 3: Query X = 4
upper_bound(4) = 4 (index of first element > 4)
count = 4
sum = prefixSum[4] = 10 (sum of 1+2+3+4)
Result: (4, 10)
Complexity Analysis
- Time Complexity: O(n log n + q log n)
- O(n log n) for sorting
- O(n) for prefix sum computation
- O(log n) per query using binary search
- Space Complexity: O(n) - For prefix sum array
Key Takeaways
- Binary Search on Sorted Data: Sorting enables efficient binary search queries
- Prefix Sum for Range Queries: Precomputing prefix sums allows O(1) sum queries
- Upper Bound: Use upper_bound to find the first element greater than target
- Trade-off: We spend O(n log n) preprocessing to get O(log n) per query
- Competitive Programming Pattern: This is a classic pattern for handling multiple queries efficiently
- Space-Time Trade-off: Using O(n) extra space reduces time from O(qn) to O(qlog n)