755. Pour Water

You are given an elevation map represents as an integer array heights where heights[i] representing the height of the terrain at index i. The width at each index is 1. You are also given two integers volume and k. volume units of water will fall at index k.

Water first drops at the index k and rests on top of the highest terrain or water at that index. Then, it flows according to the following rules:

• If the droplet would eventually fall by moving left, then move left.
• Otherwise, if the droplet would eventually fall by moving right, then move right.
• Otherwise, rise to its current position.

Here, “eventually fall” means that the droplet will eventually be at a lower level if it moves in that direction. Also, level means the height of the terrain plus any water in that column.

We can assume there is infinitely high terrain on the two sides out of bounds of the array. Also, there could not be partial water being spread out evenly on more than one grid block, and each unit of water has to be in exactly one block.

class Solution {
    int leftOf(vector<int>& A, int now) {
        while(now and A[now] == A[now - 1]) now--;
        if(now == 0) return -1;
        if(A[now] < A[now - 1]) return -1;
        int res = now;
        while(now and A[now] >= A[now - 1]) {
            if(A[now] > A[now - 1]) res = now - 1;
            now--;
        }
        return res;
    }
    int rightOf(vector<int>& A, int now) {
        while(now != A.size() - 1 and A[now] == A[now + 1]) now++;
        if(now == A.size() - 1) return -1;
        if(A[now] < A[now + 1]) return -1;
        int res = now;
        while(now != A.size() - 1 and A[now] >= A[now + 1]) {
            if(A[now] > A[now + 1]) res = now + 1;
            now++;
        }
        return res;
    }
public:
    vector<int> pourWater(vector<int>& A, int volume, int k) {
        while(volume--) {
            int l = leftOf(A, k);
            if(l != -1) {
                A[l]++;
            } else {
                int r = rightOf(A, k);
                if(r != -1) {
                    A[r]++;
                } else A[k]++;
            }
        }
        return A;
    }
};