Given a graph of $$$N$$$ nodes . Each node $$$i$$$ has value $$$A[i]$$$ . There is a directed edge from node $$$i$$$ to node $$$j$$$ if $$$A[i]\ >\ A[j]$$$ and $$$|i-j|\ \le\ K$$$ .
For every element , $$$B[i]$$$ is shortest distance to a node with prime value , if there is no path to such node $$$B[i]\ =\ 0$$$ . Find the sum of all elements of $$$B$$$ modulo $$$10^9\ +\ 7$$$ .
Constraints :
$$$1 \le K\le N\ \le 10^5\newline $$$ $$$B[i]\le\ 10^6\ \ \ \ \forall\ 1 \le i \le N$$$
Example :
- $$$N=1 \ , K=1 \ , A = [1]\newline$$$ $$$B = [0]\ , sum = 0\newline$$$
- $$$N=3 \ , K=1 \ , A = [1,2,4]\newline$$$ $$$B = [0,0,1] \ , sum = 1\newline$$$
- $$$N=5 \ , K=2 \ , A = [2,4,8,16,32]\newline$$$ $$$B = [0,1,1,2,2] \ , sum = 6\newline$$$
Note : This question was asked in an on-campus coding exam & it was completed
Process nodes from smallest $$$A[i]$$$ to largest; if $$$A[i] > A[j]$$$, the $$$j$$$'th node is unaffected by the presence or not of node $$$i$$$, as $$$j$$$ can never get to $$$i$$$. We can thus compute the final answer to such a node $$$j$$$ and then go on to $$$i$$$ later. At the start, let all $$$B[i]$$$ be $$$inf$$$.
Everytime a node is processed (imagine as if it 'began to exist from now on'), it's $$$B[i]$$$ is either $$$0$$$ if $$$A[i]$$$ is prime, or $$$1 + min(B[i-K],\cdots,B[i+K])$$$ (effectively $$$inf$$$ if all of them are $$$inf$$$).
You can do this operation for example with a minimum segment tree, in $$$O(logN)$$$.
At the end you have $$$B$$$, where $$$B[i] >= inf$$$ indicates there is no path to a prime node so print $$$0$$$.