967. Numbers With Same Consecutive Differences

Python

DFS

class Solution:
    def numsSameConsecDiff(self, n: int, k: int) -> List[int]:
        if n == 1:
            return [num for num in range(10)]

        ans = set()
        def dfs(num):
            if len(str(num)) == n:
                ans.add(num)
                return

            tail = num % 10

            for digit in [tail+k, tail-k]:
                if (digit >= 0 and digit < 10):
                    new_num = num*10 + digit
                    dfs(new_num)

        for num in range(1, 10):
            dfs(num)

        return list(ans)

Go

DFS

func numsSameConsecDiff(n int, k int) []int {
    ans := make(map[int]int)

    for num := 1; num < 10; num++ {
        dfs(n, k, num, ans)
    }

    keys := make([]int, 0, len(ans))
    for key := range(ans) {
        keys = append(keys, key)
    }

    return keys
}


func dfs(n int, k int, num int, ans map[int]int) {
    if (len(strconv.Itoa(num)) == n) {
        ans[num] = num
        return
    }

    tail := num % 10

    upper := tail + k
    lower := tail - k

    if upper < 10 {
        new_num := num*10 + upper
        dfs(n, k, new_num, ans)
    }
    if lower >= 0 {
        new_num := num*10 + lower
        dfs(n, k, new_num, ans)
    }
}