Coding questions · Amazon
Maximize the Bitwise AND of K Elements After Spending an Increment Budget
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.
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
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
kchosen elements must be atkdifferent 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^51 <= k <= len(arr)0 <= arr[i] <= 10^90 <= c <= 10^9- No element can exceed
2 * 10^9after the increases, so the answer fits in a 32-bit signed integer. Intermediate sums of increase amounts can exceed2^31 - 1, so use 64-bit integers in languages with fixed-width types.
Examples
Example 1
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
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
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.
Largest k-Element AND After Budgeted Increases
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.
Approach
Greedy over bits from high to low. Call a mask T feasible if some allowed increases and some k distinct indices give an AND that contains every bit of T. The AND contains T exactly when every chosen element ends at a supermask of T, so the cheapest way to use element x is to raise it to the smallest supermask y of T with y >= x. If x already contains T, y = x and the cost is 0. Otherwise let h be the highest bit of T that x lacks: any y > x first differs from x at some bit p where y has 1 and x has 0, and the bits above p must already contain T's bits there, which forces p >= h; p = h is valid and gives the smallest y, namely x's bits above h followed by T's bits at positions 0..h, i.e. y = (x & ~low) | (T & low) with low = 2^(h+1) - 1. Unchosen elements need no increase, so T is feasible exactly when the sum of the k smallest costs y - x is at most c. Feasibility is closed under removing bits from T, so starting from ans = 0 and trying bits 30 down to 0 (no value can exceed 2 * 10^9 < 2^31), set bit b whenever ans | 2^b is feasible. Invariant: after bit b is processed, ans equals the optimum restricted to bits >= b. If the optimum has bit b, ans | 2^b is a submask of it and therefore feasible; if it lacks bit b, a feasible ans | 2^b would give an AND larger than the optimum, a contradiction. Edge cases: k = 1 yields max(arr) + c; c = 0 reduces to the best AND of k existing elements; raising an element across a power of two clears its lower bits, so 1 is best raised to 2 rather than 3 in [1, 2], k = 2, c = 1; an element larger than the target can still lack its bits (9 lacks the bits of 6); and the sum of k costs can reach about 2 * 10^14, so it is accumulated in 64-bit integers.
- Time complexity: O(31 * n log n)
- Space complexity: O(n)