Computing optimal price for data plan

Question

I am not sure if this is the right place to ask this question since it is actually an algorithmic question but I think that it requires some math technique as well:

A data plan has a cost of $ X $. We are given list $G$ containing $N$ (unsorted) integers, $Y$ (1 $ \le Y \le $ 1e6). A cost function $C(i,X)$ is defined as follows:

$C(i,X)$ represents the cost of using a data plan, of price $X$, up to the $i$ th integer in the list, $ 1 \le i \le N $

$\forall Y \in G,$ \begin{equation} C(i,X) = \begin{cases} Y \gt X \implies C(i-1,X) + X + \\ (Y - X) \cdot (Y - X) \\ Y \le X \implies C(i-1,X) + X \end{cases} \end{equation}

Please find a value of $ X $ so that $C(N,X)$ is minimized.

On first thought, a $O(N^2)$ method is trivial but is very inefficient for large input (i.e. $N=$ 2e5) . For any reasonably large input in this problem, an efficient algorithm should have a complexity of $ O(NlogN) $ or better.

My approach is to use binary search. As a precursor, I sort the array. Then I imagine drawing a line in a bar chart based on the integers in the array - the line divides the chart into 2 portions: each of the elements in the bottom half $\le X $ and those in the other half $\gt X$.

Let total cost incurred by elements greater than $X$ be $H$ and the total cost incurred by the rest be $L$. Let $BOT$ represent the lower bound for my guess and $TOP$ represent the upper bound for my guess.

Here is the pseudo-code for my logic:

WHILE BOT less than or equals to TOP:
  DO
    SET X = (BOT+TOP)/2;
    // Compute H and L based on X
    // ...
    SET ANS = MIN(ans,H+L)
    IF H equals to L:
      EXIT LOOP
    ELSE IF H > L:
      SET BOT = X + 1
    ELSE:
      SET TOP = X - 1
  DONE

PRINT ANS

Example:

1 2 9 5 7 6 3 4 2 2

Answer should be 70 because we can choose 6.

Sorting the array:

Elements: 1 2 2 2 3 4 5 6 7 9
Cost:     6 6 6 6 6 6 6 6 7 15

Hence, 6(8) + 7 + 15 = 70

I don't have any formal proof for my solution and it turns out that my logic is wrong. Could anyone please advise me of a better solution?

Answer 1

It might be useful to rewrite the formula for$C(X)$.Let the ordered values be $Y_1\leq Y_2\leq \dots \leq Y_N$.

If $k$ is the first value such that $Y_k\geq X$ then $C(X)=NX+\sum_{k}^N (Y_k-X)^2$. Clearly $C(x)-C(x-1)$ is a decreasing function. So we just have to find the greatest value $x$ such that $C(x)-C(x-1)$ is still positive, and this can be done with a simple binary search, the complexity is $\mathcal O(n)$ (the number of steps in the binary search is at most $\log(10^6)$, so we have to take into account a constant of about $20$).

Some c++ code:

#include 
typedef unsigned long long lli;

int inf=1e6;
int y[1000010];
int n;

lli cost(int x){
    lli res=0;
    for(int i=0;ix) res+=(y[i]-x)*(y[i]-x);
    }
    return(res);
}

int good(int x){
    if(cost(x+1)<=cost(x)) return(1);
    return(0);
}

int main(){
    scanf("%d",&n);
    for(int i=0;i