Maximize the Bitwise AND of K Elements After Spending an Increment Budget
Company: Amazon
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: hard
Interview Round: Online Assessment
You are given an array `arr` of non-negative integers, an integer `k` and a non-negative integer budget `c`.
First, you may increase elements of `arr`. Each increase adds a non-negative integer amount to one element, and the total added across all elements must not exceed `c`. You may spread the budget over several elements, put all of it on one element, or leave some of it unused. Elements can only be increased, never decreased.
Then you choose exactly `k` elements at distinct indices (they do not need to be contiguous) and take the bitwise AND of their values after the increases.
Return the largest bitwise AND you can obtain.
### Function Signature
```python
def max_and_after_increments(arr: list[int], k: int, c: int) -> int:
```
### Rules
- Every increase is a non-negative integer, and the sum of all increases is at most `c`.
- The `k` chosen elements must be at `k` different indices. Equal values at different indices may both be chosen.
- Increasing an element that is not chosen is allowed, but it simply wastes budget.
- With `k = 1`, the bitwise AND of a single element is that element's value.
- Return the maximum achievable AND value. It is unique even when several choices of increases and indices reach it.
### Constraints
- `1 <= len(arr) <= 10^5`
- `1 <= k <= len(arr)`
- `0 <= arr[i] <= 10^9`
- `0 <= c <= 10^9`
- No element can exceed `2 * 10^9` after the increases, so the answer fits in a 32-bit signed integer. Intermediate sums of increase amounts can exceed `2^31 - 1`, so use 64-bit integers in languages with fixed-width types.
### Examples
**Example 1**
```text
Input: arr = [5, 4, 1, 7, 2], k = 3, c = 3
Output: 6
```
Increase `4` by 2 and `5` by 1 (3 in total), then choose `6`, `6` and `7`: `6 & 6 & 7 = 6`. No allocation of at most 3 lets any three elements reach an AND of 7 or more.
**Example 2**
```text
Input: arr = [12, 10, 14, 7], k = 2, c = 0
Output: 12
```
With no budget, the best pair is `12` and `14`: `12 & 14 = 12`.
**Example 3**
```text
Input: arr = [1, 2], k = 2, c = 1
Output: 2
```
Adding 1 to the first element gives `[2, 2]` and `2 & 2 = 2`. Adding it to the second element gives `1 & 3 = 1`, and adding nothing gives `1 & 2 = 0`.
Overview: Given an array, a count k and a total budget of increments, raise elements so that the bitwise AND of some k chosen elements is as large as possible, and return that value. The problem tests bit manipulation, careful accounting of a shared budget, and an efficient solution for inputs of up to 100,000 elements.
Read the full Amazon Software Engineer interview experience this question came from
You are given an array `arr` of non-negative integers, an integer `k`, and a non-negative integer budget `c`.
First, you may increase elements of `arr`. Each increase adds a non-negative integer amount to one element, and the total amount added across all elements must not exceed `c`. You may spread the budget over several elements, put all of it on one element, or leave some of it unused. Elements can only be increased, never decreased.
Then you choose exactly `k` elements at distinct indices (they do not need to be contiguous) and take the bitwise AND of their values after the increases.
Return the largest bitwise AND you can obtain.
Rules:
- Every increase is a non-negative integer, and the sum of all increases is at most `c`.
- The `k` chosen elements must be at `k` different indices. Equal values at different indices may both be chosen.
- Increasing an element that is not chosen is allowed, but it simply wastes budget.
- With `k = 1`, the bitwise AND of a single element is that element's value.
- Return the maximum achievable AND value as a single integer. It is unique even when several choices of increases and indices reach it.
Example 1:
Input: arr = [5, 4, 1, 7, 2], k = 3, c = 3
Output: 6
Explanation: Increase 4 by 2 and 5 by 1 (3 in total), then choose 6, 6 and 7: 6 & 6 & 7 = 6. No allocation of at most 3 lets any three elements reach an AND of 7 or more.
Example 2:
Input: arr = [1, 2], k = 2, c = 1
Output: 2
Explanation: Adding 1 to the first element gives [2, 2] and 2 & 2 = 2. Adding it to the second element gives 1 & 3 = 1, and adding nothing gives 1 & 2 = 0.
Constraints:
- 1 <= len(arr) <= 10^5
- 1 <= k <= len(arr)
- 0 <= arr[i] <= 10^9
- 0 <= c <= 10^9
- No element can exceed 2 * 10^9 after the increases, so the answer fits in a 32-bit signed integer. Intermediate sums of increase amounts can exceed 2^31 - 1, so use 64-bit integers in languages with fixed-width types (Java long, C++ long long).
The inputs and the returned value all fit in a signed 32-bit integer; only sums of increase amounts may exceed 2^31 - 1.
Constraints
- 1 <= len(arr) <= 10^5
- 1 <= k <= len(arr)
- 0 <= arr[i] <= 10^9
- 0 <= c <= 10^9
- No element can exceed 2 * 10^9 after the increases, so the answer fits in a 32-bit signed integer. Intermediate sums of increase amounts can exceed 2^31 - 1, so use 64-bit integers in languages with fixed-width types.
Examples
Input: ([5, 4, 1, 7, 2], 3, 3)
Expected Output: 6
Explanation: Source Example 1: raise 5 by 1 and 4 by 2 and pick 6, 6, 7; making three elements contain 7 costs at least 5.
Input: ([12, 10, 14, 7], 2, 0)
Expected Output: 12
Explanation: Source Example 2: no budget, so the best pair is 12 & 14 = 12.
Hints
- With k = 1 the answer is simply the largest element plus c; the difficulty is in sharing the budget among k chosen elements.
- Adding to an element can clear its lower bits: in Example 2, raising 1 to 2 loses bit 0 but lets both elements share bit 1.
- A single higher bit in the AND is worth more than all lower bits combined.
Community answers
Answer by patisid
import java.util.*;
class Solution {
public static int maxAndAfterIncrements(int[] arr, int k, long c) {
int n = arr.length;
long ans = 0;
long[] cost = new long[n];
for (int b = 30; b >= 0; b--) { // arr[i] + c <= 2e9 < 2^31
long T = ans | (1L << b);
for (int i = 0; i < n; i++) cost[i] = costTo(arr[i], T);
Arrays.sort(cost);
long sum = 0;
for (int i = 0; i < k && sum <= c; i++) sum += cost[i];
if (sum <= c) ans = T;
}
return (int) ans;
}
// min (y - x) with y >= x and (y & T) == T
private static long costTo(long x, long T) {
long missing = T & ~x;
if (missing == 0) return 0;
int p = 63 - Long.numberOfLeadingZeros(missing); // highest bit T needs that x lacks
long y = (((x >> p) | 1) << p) | T;
return y - x;
}
}