Search code examples
c++algorithmrecursionmemoization

Why is memoization giving the wrong output?


I am tryna solve a question using recursion+memoization. It is just modified Fibonacci with an additional start+3 step that can be done only K times. This is the recursive code I came up with:

#include <iostream>
using namespace std;

int helper(int start, int N, int K) {
    // cout<<start<<" "<<N<<" "<<K<<"\n";
    if(start>N) return 0;
    if(start==N) return 1;
    
    int ans=0;
    ans=helper(start+1, N, K) + helper(start+2, N, K);
    if(K>0) {
        K--;
        ans=ans+helper(start+3, N, K);
    }
    
    return ans;
}

int main() {
    int T;
    cin>>T;
    while(T--) {
        int N, K;
        cin>>N>>K;
        cout<<helper(0, N, K)<<"\n";
    
    }
    
    return 0;
}

Just memoizing it and taking modulo 10^9+7 (the problem requires me to do it), I have:

#include <iostream>
#include <vector>
using namespace std;

vector<long long int> dp;
const unsigned int M = 1000000007;

int helper(long long int start, long long int N, int K) {
    // cout<<start<<" "<<N<<" "<<K<<"\n";
    if(start>N) return 0;
    if(start==N) return 1;
    if(dp[start]!=0) {
        return dp[start];
    }
    
    int ans=0;
    ans=(helper(start+1, N, K)%M + helper(start+2, N, K)%M)%M;
    if(K>0) {
        K--;
        ans=(ans%M+helper(start+3, N, K)%M)%M;
    }
    
    return dp[start]=ans;
}

int main() {
    int T;
    cin>>T;
    while(T--) {
        long long int N;
        int K;
        cin>>N>>K;
        dp.clear();
        dp.resize(N+5, 0);
        cout<<helper(0, N, K)<<"\n";
    }
    
    return 0;
}

The code is exactly the same, but for memoization and modulo. When I run it on the following input:

1
7 1

I get 41 in the first case, while 44 in the second one. Obviously, I debugged and I expected some memoization or modulo issues. However, I noticed that some of the evaluations are no longer being computed, by which I concluded some issue with the recursive calls. Please note the diff of the calls here and working code snippets here and here.

Could someone please point out what I am missing?


Solution

  • You're caching results in db based entirely on start, and ignoring K entirely. But helper(0, 7, 0) and helper(0,7,7) need to cache entirely different numbers at each index, based on whatever K is remaining. So you need the cache index to also be based on K.

    I see three ways to optimize:

    • Have a cache of caches. db[start] would be a std::vector, and you'd use index K in that vector to access your item. This results in using significantly more memory.
    • Abandon std::vector and use a std::unordered_map, where your key is a pair of (start, K). (Possibly std::pair<int, int>, but better to use a struct with better names.) This should use drastically less memory, with little lost in speed.
    • Only cache values when K==0. This doesn't increase your cache size, but you also lose a lot of the caching benefits when K is large.