Problem

You are given a non-negative integer k. There exists a staircase with an infinite number of stairs, with the lowest stair numbered 0.

Alice has an integer jump, with an initial value of 0. She starts on stair 1 and wants to reach stair k using any number of operations. If she is on stair i, in one operation she can:

  • Go down to stair i - 1. This operation cannot be used consecutively or on stair 0.
  • Go up to stair i + 2jump. And then, jump becomes jump + 1.

Return the total number of ways Alice can reach stair k.

Note that it is possible that Alice reaches the stair k, and performs some operations to reach the stair k again.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


Input: k = 0

Output: 2

Explanation:

The 2 possible ways of reaching stair 0 are:

  * Alice starts at stair 1. 
* Using an operation of the first type, she goes down 1 stair to reach stair 0.
  * Alice starts at stair 1. 
* Using an operation of the first type, she goes down 1 stair to reach stair 0.
* Using an operation of the second type, she goes up 20 stairs to reach stair 1.
* Using an operation of the first type, she goes down 1 stair to reach stair 0.

Example 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


Input: k = 1

Output: 4

Explanation:

The 4 possible ways of reaching stair 1 are:

  * Alice starts at stair 1. Alice is at stair 1.
  * Alice starts at stair 1. 
* Using an operation of the first type, she goes down 1 stair to reach stair 0.
* Using an operation of the second type, she goes up 20 stairs to reach stair 1.
  * Alice starts at stair 1. 
* Using an operation of the second type, she goes up 20 stairs to reach stair 2.
* Using an operation of the first type, she goes down 1 stair to reach stair 1.
  * Alice starts at stair 1. 
* Using an operation of the first type, she goes down 1 stair to reach stair 0.
* Using an operation of the second type, she goes up 20 stairs to reach stair 1.
* Using an operation of the first type, she goes down 1 stair to reach stair 0.
* Using an operation of the second type, she goes up 21 stairs to reach stair 2.
* Using an operation of the first type, she goes down 1 stair to reach stair 1.

Constraints

  • 0 <= k <= 10^9

Solution

Method 1 – Dynamic Programming with Memoization

Intuition

At each step, Alice can either go up (with a jump that increases each time) or go down (but not consecutively and not from stair 0). The state is defined by the current stair, current jump, and whether the last move was down. We use memoization to avoid recomputation.

Approach

  1. Define a recursive function dp(i, jump, last_down) where:
    • i is the current stair,
    • jump is the current jump value,
    • last_down is a boolean indicating if the last move was down.
  2. Base case: if i == k, count as a valid way.
  3. For each state:
    • If last move was not down and i > 0, try going down to i-1 (set last_down=True).
    • Always try going up to i + 2*jump (set last_down=False, jump+1).
  4. Use memoization to cache results for each state.
  5. Return the total number of ways from the initial state (i=1, jump=0, last_down=False).

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution {
public:
    int numberOfWays(int k) {
        map<tuple<int,int,bool>, int> memo;
        function<int(int,int,bool)> dp = [&](int i, int jump, bool last_down) -> int {
            if (i == k) return 1;
            auto key = make_tuple(i, jump, last_down);
            if (memo.count(key)) return memo[key];
            int res = 0;
            if (!last_down && i > 0) res += dp(i-1, jump, true);
            res += dp(i + 2*jump, jump+1, false);
            return memo[key] = res;
        };
        return dp(1, 0, false);
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

func numberOfWays(k int) int {
    memo := map[[3]int]int{}
    var dp func(i, jump, lastDown int) int
    dp = func(i, jump, lastDown int) int {
        if i == k { return 1 }
        key := [3]int{i, jump, lastDown}
        if v, ok := memo[key]; ok { return v }
        res := 0
        if lastDown == 0 && i > 0 {
            res += dp(i-1, jump, 1)
        }
        res += dp(i+2*jump, jump+1, 0)
        memo[key] = res
        return res
    }
    return dp(1, 0, 0)
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution {
    public int numberOfWays(int k) {
        Map<String, Integer> memo = new HashMap<>();
        return dp(1, 0, false, k, memo);
    }
    private int dp(int i, int jump, boolean lastDown, int k, Map<String, Integer> memo) {
        if (i == k) return 1;
        String key = i + "," + jump + "," + lastDown;
        if (memo.containsKey(key)) return memo.get(key);
        int res = 0;
        if (!lastDown && i > 0) res += dp(i-1, jump, true, k, memo);
        res += dp(i+2*jump, jump+1, false, k, memo);
        memo.put(key, res);
        return res;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution {
    fun numberOfWays(k: Int): Int {
        val memo = mutableMapOf<Triple<Int,Int,Boolean>, Int>()
        fun dp(i: Int, jump: Int, lastDown: Boolean): Int {
            if (i == k) return 1
            val key = Triple(i, jump, lastDown)
            if (key in memo) return memo[key]!!
            var res = 0
            if (!lastDown && i > 0) res += dp(i-1, jump, true)
            res += dp(i+2*jump, jump+1, false)
            memo[key] = res
            return res
        }
        return dp(1, 0, false)
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class Solution:
    def numberOfWays(self, k: int) -> int:
        from functools import lru_cache
        @lru_cache(None)
        def dp(i: int, jump: int, last_down: bool) -> int:
            if i == k:
                return 1
            res = 0
            if not last_down and i > 0:
                res += dp(i-1, jump, True)
            res += dp(i+2*jump, jump+1, False)
            return res
        return dp(1, 0, False)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

impl Solution {
    pub fn number_of_ways(k: i32) -> i32 {
        use std::collections::HashMap;
        fn dp(i: i32, jump: i32, last_down: bool, k: i32, memo: &mut HashMap<(i32,i32,bool), i32>) -> i32 {
            if i == k { return 1; }
            if let Some(&v) = memo.get(&(i, jump, last_down)) { return v; }
            let mut res = 0;
            if !last_down && i > 0 {
                res += dp(i-1, jump, true, k, memo);
            }
            res += dp(i+2*jump, jump+1, false, k, memo);
            memo.insert((i, jump, last_down), res);
            res
        }
        let mut memo = HashMap::new();
        dp(1, 0, false, k, &mut memo)
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution {
    numberOfWays(k: number): number {
        const memo = new Map<string, number>();
        function dp(i: number, jump: number, lastDown: boolean): number {
            if (i === k) return 1;
            const key = `${i},${jump},${lastDown}`;
            if (memo.has(key)) return memo.get(key)!;
            let res = 0;
            if (!lastDown && i > 0) res += dp(i-1, jump, true);
            res += dp(i+2*jump, jump+1, false);
            memo.set(key, res);
            return res;
        }
        return dp(1, 0, false);
    }
}

Complexity

  • ⏰ Time complexity: O(k^2), since the state space is O(k^2) (i and jump) and each state is computed once.
  • 🧺 Space complexity: O(k^2), for memoization cache.