Construct the Array from HackerRank

Question

This is my answer for https://www.hackerrank.com/challenges/construct-the-array/problem (Passes all the test cases)

Problem

Your goal is to find the number of ways to construct an array such that consecutive positions contain different values.

Specifically, we want to construct an array with \$n\$ elements such that each element [contains a value] between \1ドル\$ and \$k\$, inclusive. We also want the first and last elements of the array to be \1ドル\$ and \$x\$. Given \$n\$, \$k\$ and \$x\$, find the number of ways to construct such an array.

Since the answer may be large, only find answer modulo \10ドル^9 + 7\$

For example, for \$n=4\$, \$k=3\$ and \$x=2\$ there are \3ドル\$ ways

Solution

public class Solution {
 public static final long MOD = 1_000_000_007L;
 public static long countArray(int n, int k, int x) {
 if (n < 3 || k < 2) return 0;
 if (n == 3 && x == 1) {
 return k - 1;
 }
 if (n == 3) {
 return k - 2;
 }
 long countOne = k - 1, countOther = k - 2, temp;
 for (int i = 4; i < n; i++) {
 temp = countOther * (k - 1) % MOD;
 countOther = (countOne % MOD + countOther * (k - 2) % MOD) % MOD;
 countOne = temp;
 }
 if (x == 1) {
 return countOther * (k - 1) % MOD;
 }
 return (countOne % MOD + countOther * (k - 2) % MOD) % MOD;
 }
}

Answer 1

"Impossible" check

Per task description, k >= 2 and n >= 3. But nevertheless you're performing this "impossible" check:

if (n < 3 || k < 2) return 0;

Redundant check

I don't understand why you're checking cases n == 3 separately. You could remove these 2 checks and have this case handled in your for loop quite easily as below (please note that the values of i, countOne and countOther have changed):

long countOne = 0, countOther = 1, temp;
for (int i = 3; i < n; i++) {
 temp = countOther * (k - 1) % MOD;
 countOther = (countOne % MOD + countOther * (k - 2) % MOD) % MOD;
 countOne = temp;
} 

Variable naming

countOne or countOther are not a good way to call your variables - their names don't convey their purpose. I'll leave it up to you to come up with better names.

Excessive modulo operation

You don't have to perform a modulo operation more times than is needed. Given you task constraints, your can perform a modulo operation twice per iteration as below:

temp = countOther * (k - 1) % MOD;
countOther = (countOne + countOther * (k - 2)) % MOD;
countOne = temp;

Reduce scope of temp

As @Barmar correctly pointed out, your temp variable should be local to the for loop. It's not needed anywhere out of it.

Simplify return statements

You don't have to calculate countOne or countOther once you leave the for loop. You can do it all there, just fix your loop condition from i < n to i <= n and swap countOne with countOther in your return statements.

for (int i = 3; i <= n; i++) {
 long temp = countOther * (k - 1) % MOD;
 countOther = (countOne + countOther * (k - 2)) % MOD;
 countOne = temp;
}
if (x == 1) {
 return countOne;
}
return countOther;

Final Code

public static long countArray(int n, int k, int x) {
 long countOne = 0, countOther = 1;
 for (int i = 3; i <= n; i++) {
 long temp = countOther * (k - 1) % MOD;
 countOther = (countOne + countOther * (k - 2)) % MOD;
 countOne = temp;
 }
 return x == 1 ? countOne : countOther;
}

• JaDogg: I chose this answer becuase I like how you simplified my code. Thanks :)

Answer 2

I suggest that you extract the similar evaluations in methods.

 private static long countFirst(int k, long countOther, int i) {
 return countOther * (k - i) % MOD;
 }
 private static long countNext(int k, long countOne, long countOther) {
 return (countOne % MOD + countFirst(k, countOther, 2)) % MOD;
 }

//[...]
 long countOne = k - 1;
 long countOther = k - 2;
 long temp;
 for (int i = 4; i < n; i++) {
 temp = countFirst(k, countOther, 1);
 countOther = countNext(k, countOne, countOther);
 countOne = temp;
 }
 if (x == 1) {
 return countFirst(k, countOther, 1);
 }
 return countNext(k, countOne, countOther);
//[...]
```

Answer 3

You are checking conditions multiple times, increasing the complexity of the code. For example:

 if (n == 3 && x == 1) {
 return k - 1;
 }
 if (n == 3) {
 return k - 2;
 }

checks n == 3 twice. A nested if statement would be simpler & clearer:

 if (n == 3) {
 if (x == 1) {
 return k - 1;
 } else {
 return k - 2;
 }
 }

---

Despite passing all the tests, this code may give the wrong result if n = 3 and k > MOD, unless the problem imposed a limit on k not reflected in the above problem description.

• JaDogg: You are correct k isn't larger than MOD :)