House Robber III dynamiic programming memoization with arrays and tree payloads

house-robber-iii-array.cpp

struct TreeNode {
   int val;
   TreeNode *left;
   TreeNode *right;
   TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

int kaux;

struct IndexedTreeNode {
    int val, k;
    IndexedTreeNode *left, *right;
    IndexedTreeNode(TreeNode *n, int k): k(k), val(n->val), left(NULL), right(NULL) { } // make sure n is not null
};

int size(TreeNode* node, IndexedTreeNode *knode) {
    if (!node) return 0;
    int k = knode->k;
    
    if (node->left) knode->left = new IndexedTreeNode(node->left, kaux++);
    if (node->right) knode->right = new IndexedTreeNode(node->right, kaux++);
    
    return 1 + size(node->left, knode->left) + size(node->right, knode->right);
}

int dyn(IndexedTreeNode *knode, int *d) {
    return knode ? d[knode->k] : 0;
}

int compute(IndexedTreeNode *knode, int *dk, int *dkb, bool *s) {
    if (!knode) return 0;
    int k = knode->k;
    
    if (!s[k]) {
        compute(knode->left, dk, dkb, s);
        compute(knode->right, dk, dkb, s);
        
        dk[k] = knode->val + dyn(knode->left, dkb) + dyn(knode->right, dkb);
        dkb[k] = max( dyn(knode->left, dk), dyn(knode->left, dkb) ) + max( dyn(knode->right, dk), dyn(knode->right, dkb) );
        s[k] = true;
    }
    
    return max(dk[k], dkb[k]);
}

int rob(TreeNode* root) {
    if (!root) return 0;
    
    kaux = 0;
    IndexedTreeNode *kroot = new IndexedTreeNode(root, kaux++);

    int n = size(root, kroot);
    
    int dk[n], dkb[n];
    bool s[n]; for (int i = 0; i < n; ++i) s[i] = false;
    
    return compute(kroot, dk, dkb, s);
}

house-robber-iii-payload.cpp

struct TreeNode {
   int val;
   TreeNode *left, *right;
   TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

struct Payload {
    int dk, dkb;
};

void compute(TreeNode *n, Payload *p) {
    if (!n) return;
    
    Payload *pl = 0, *pr = 0;
    TreeNode *nl = n->left, *nr = n->right;
    if (nl) pl = new Payload(), compute(nl, pl);
    if (nr) pr = new Payload(), compute(nr, pr);
    
    p->dk  = n->val + (pl ? pl->dkb : 0) + (pr ? pr->dkb : 0);
    p->dkb = (pl ? max( pl->dk, pl->dkb ) : 0) + (pr ? max( pr->dk, pr->dkb ) : 0);
    
    if (pl) delete pl;
    if (pr) delete pr;
}

int rob(TreeNode* root) {
    if (!root) return 0;
    Payload *proot = new Payload();
    
    compute(root, proot);
    return max(proot->dk, proot->dkb);
}