Search code examples
arrayssortingtime-complexityincrementdecrement

Increment or decrement array elements by 1


I need some help solving the following competitive programming question. Given array A of size N .now you can make K operations where in each operation you can increase or decrease an element in array by 1. Now beauty of array is frequency of most frequent element in the array. Now output max beauty that can be achieved in K operations. This function takes the following 3 parameters and returns the required answer N- Represents the size of array nums K- Represents the integer K nums- Represents the elements of array nums Input format for custom testing-

The first line contains T, which represents the number of test cases. For each test case- The first line contains N denoting the size of array nums The second line contains integer K. The third line contains N space-separated integers of array A

`Constraints:`1\<T\<10
1\<N\<1e5
0\<K\<1e18
\-1e9\<Ai\<1e9
\`

I tried the following approach , by sorting the array, then finding consecutive differences between them, to get max elements equal in k operations , I need to choose elements from index i to j such that no of operations are min and beauty is maximum.


Solution

  • vector<int> preSum;
    
    void calcPresum(vector<int>& v) {
        int n = v.size();
        preSum.clear(); preSum.resize(n, v[0]);
        for(int i = 1; i < n; i++) {
            preSum[i] = preSum[i - 1] + v[i];
        }
    }
    
    int rangeSum(int i, int j) {
        if(i == 0) return preSum[j];
        return preSum[j] - preSum[i - 1];
    }
    
    int getCost(vector<int>& v, int st, int en, int mid) {
        int costLeft = (v[mid] * (mid - st + 1)) - rangeSum(st, mid);
        int costRight = rangeSum(mid, en) - (v[mid] * (en - mid + 1));
        return costLeft + costRight;
    }
    
    int costEq(vector<int>& v, int k) {
        int n = v.size(), ptr1 = 0, ptr2 = k - 1, cost = INT_MAX;
        for(; ptr2 < n; ptr2++, ptr1++) {
            // decide what is min cost in range ptr1 - ptr2
            // if odd elements median will give min cost
            if((ptr2 - ptr1 + 1) % 2 == 1) {
                cost = min(cost, getCost(v, ptr1, ptr2, (ptr1 + ptr2) / 2));
            }
            // if even elements check which median will give min cost
            else {
                cost = min({cost, getCost(v, ptr1, ptr2, (ptr1 + ptr2) / 2), getCost(v, ptr1, ptr2, (ptr1 + ptr2) / 2 + 1)});
            }
        }
        return cost;
    }
    
    int solve(vector<int> v, int k) {
        int n = v.size();
        int ptr1 = 1, ptr2 = n;
        calcPresum(v);
        while(ptr1 < ptr2) {
            int mid = ptr1 + (ptr2 - ptr1 + 1) / 2;
            if(costEq(v, mid) <= k) {
                ptr1 = mid;
            } else {
                ptr2 = mid - 1;
            }
        }
        return ptr1;
    }
    
    int main() {
        // median gets the minimal cost for making all elements equal
        cout << solve({1, 4, 5, 7}, 4) << endl;
        cout << solve({1, 4, 4, 4, 4, 4, 4, 5, 7}, 4) << endl;
        cout << solve({1, 4, 4, 4, 4, 4, 4, 5, 7}, 2) << endl;
    }