LeetCode-in-Java

3017. Count the Number of Houses at a Certain Distance II

Hard

You are given three positive integers n, x, and y.

In a city, there exist houses numbered 1 to n connected by n streets. There is a street connecting the house numbered i with the house numbered i + 1 for all 1 <= i <= n - 1 . An additional street connects the house numbered x with the house numbered y.

For each k, such that 1 <= k <= n, you need to find the number of pairs of houses (house1, house2) such that the minimum number of streets that need to be traveled to reach house2 from house1 is k.

Return a 1-indexed array result of length n where result[k] represents the total number of pairs of houses such that the minimum streets required to reach one house from the other is k.

Note that x and y can be equal.

Example 1:

Input: n = 3, x = 1, y = 3

Output: [6,0,0]

Explanation: Let’s look at each pair of houses:

Example 2:

Input: n = 5, x = 2, y = 4

Output: [10,8,2,0,0]

Explanation: For each distance k the pairs are:

Example 3:

Input: n = 4, x = 1, y = 1

Output: [6,4,2,0]

Explanation: For each distance k the pairs are:

Constraints:

Solution

public class Solution {
    public long[] countOfPairs(int n, int x, int y) {
        long[] result = new long[n];
        int leftCount = Math.min(x, y) - 1;
        int rightCount = n - Math.max(x, y);
        int circleCount = n - leftCount - rightCount;
        circleInternal(circleCount, result);
        lineToCircle(leftCount, circleCount, result);
        lineToCircle(rightCount, circleCount, result);
        lineToLine(leftCount, rightCount, x == y ? 1 : 2, result);
        lineInternal(leftCount, result);
        lineInternal(rightCount, result);
        return result;
    }

    private void lineToCircle(int lineCount, int circleCount, long[] curRes) {
        int circleLen = circleCount / 2 + 1;
        int curModifier = 0;
        for (int i = 1; i < circleLen + lineCount; ++i) {
            if (i <= Math.min(lineCount, circleLen)) {
                curModifier += 4;
            } else if (i > Math.max(lineCount, circleLen)) {
                curModifier -= 4;
            }
            curRes[i - 1] += curModifier;
            if (i <= lineCount) {
                curRes[i - 1] -= 2;
            }
            if (i >= circleLen && circleCount % 2 == 0) {
                curRes[i - 1] -= 2;
            }
        }
    }

    private void lineToLine(int lineCount1, int lineCount2, int initialDis, long[] curRes) {
        int curModifier = 0;
        for (int i = 1; i < lineCount1 + lineCount2; ++i) {
            if (i <= Math.min(lineCount1, lineCount2)) {
                curModifier += 2;
            } else if (i > Math.max(lineCount1, lineCount2)) {
                curModifier -= 2;
            }
            curRes[i - 1 + initialDis] += curModifier;
        }
    }

    private void lineInternal(int lineCount, long[] curRes) {
        for (int i = 1; i < lineCount; ++i) {
            curRes[i - 1] += (lineCount - i) * 2L;
        }
    }

    private void circleInternal(int circleCount, long[] curRes) {
        for (int i = 0; i < circleCount / 2; ++i) {
            if (circleCount % 2 == 0 && i + 1 == circleCount / 2) {
                curRes[i] += circleCount;
            } else {
                curRes[i] += circleCount * 2L;
            }
        }
    }
}