import math def solution(n): memo = {} return get_staircase_count(n, n-1, memo) def get_staircase_count(blocks, cap, memo): if blocks <= 0: return 0 if cap <= 0: return 0 if blocks == 1: return 1 if cap >= blocks: return 1 + get_staircase_count(blocks, blocks-1, memo) if blocks not in memo: memo[blocks] = {} if cap in memo[blocks]: return memo[blocks][cap] count = 0 for height in range(1, cap+1): count += get_staircase_count(blocks-height, height-1, memo) memo[blocks][cap] = count return count tests = [ [1, 0], [2, 0], [3, 1], [4, 1], [5, 2], [6, 3], [200, 487067745] ] for test in tests: result = solution(test[0]) print(test[0], result == test[1], result, test[1]) print ""