diff --git a/library/lms.c b/library/lms.c
index 46ea567..78c7d26 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -65,10 +65,9 @@
 
 /* Currently only support H=10 */
 #define H_TREE_HEIGHT_MAX                  10
-#define MERKLE_TREE_NODE_AM_MAX            (1u << (H_TREE_HEIGHT_MAX + 1u))
-#define MERKLE_TREE_NODE_AM(type)          (1u << (MBEDTLS_LMS_H_TREE_HEIGHT(type) + 1u))
-#define MERKLE_TREE_LEAF_NODE_AM(type)     (1u << MBEDTLS_LMS_H_TREE_HEIGHT(type))
-#define MERKLE_TREE_INTERNAL_NODE_AM(type) (1u << MBEDTLS_LMS_H_TREE_HEIGHT(type))
+#define MERKLE_TREE_NODE_AM(type)          ((size_t) 1 << (MBEDTLS_LMS_H_TREE_HEIGHT(type) + 1u))
+#define MERKLE_TREE_LEAF_NODE_AM(type)     ((size_t) 1 << MBEDTLS_LMS_H_TREE_HEIGHT(type))
+#define MERKLE_TREE_INTERNAL_NODE_AM(type) ((size_t) 1 << MBEDTLS_LMS_H_TREE_HEIGHT(type))
 
 #define D_CONST_LEN           (2)
 static const unsigned char D_LEAF_CONSTANT_BYTES[D_CONST_LEN] = {0x82, 0x82};
@@ -499,13 +498,21 @@
                             unsigned int leaf_node_id,
                             unsigned char *path )
 {
-    unsigned char tree[MERKLE_TREE_NODE_AM_MAX][MBEDTLS_LMS_M_NODE_BYTES_MAX];
+    const size_t node_bytes = MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type);
     unsigned int curr_node_id = leaf_node_id;
     unsigned int adjacent_node_id;
+    unsigned char *tree = NULL;
     unsigned int height;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    ret = calculate_merkle_tree( ctx, ( unsigned char * )tree );
+    tree = mbedtls_calloc( MERKLE_TREE_NODE_AM(ctx->params.type),
+        node_bytes );
+    if ( tree == NULL )
+    {
+        return MBEDTLS_ERR_LMS_ALLOC_FAILED;
+    }
+
+    ret = calculate_merkle_tree( ctx, tree );
     if( ret != 0 )
     {
         goto exit;
@@ -516,9 +523,8 @@
     {
         adjacent_node_id = curr_node_id ^ 1;
 
-        memcpy( &path[height * MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type)],
-                &tree[adjacent_node_id],
-                MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) );
+        memcpy( &path[height * node_bytes],
+                &tree[adjacent_node_id * node_bytes], node_bytes );
 
         curr_node_id >>=1;
     }
@@ -526,7 +532,9 @@
     ret = 0;
 
 exit:
-    mbedtls_platform_zeroize( tree, sizeof( tree ) );
+    mbedtls_platform_zeroize( tree, node_bytes *
+        MERKLE_TREE_NODE_AM(ctx->params.type) );
+    mbedtls_free ( tree );
 
     return( ret );
 }
@@ -659,8 +667,9 @@
 int mbedtls_lms_calculate_public_key( mbedtls_lms_public_t *ctx,
                                       const mbedtls_lms_private_t *priv_ctx )
 {
-    unsigned char tree[MERKLE_TREE_NODE_AM_MAX][MBEDTLS_LMS_M_NODE_BYTES_MAX];
+    const size_t node_bytes = MBEDTLS_LMS_M_NODE_BYTES(priv_ctx->params.type);
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    unsigned char *tree = NULL;
 
     if( ! priv_ctx->have_private_key )
     {
@@ -679,25 +688,33 @@
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
+    tree = mbedtls_calloc( MERKLE_TREE_NODE_AM(priv_ctx->params.type),
+        node_bytes );
+    if ( tree == NULL )
+    {
+        return MBEDTLS_ERR_LMS_ALLOC_FAILED;
+    }
+
     memcpy( &ctx->params, &priv_ctx->params,
             sizeof( mbedtls_lmots_parameters_t ) );
 
-    ret = calculate_merkle_tree( priv_ctx, ( unsigned char * )tree );
+    ret = calculate_merkle_tree( priv_ctx, tree );
     if( ret != 0 )
     {
         goto exit;
     }
 
     /* Root node is always at position 1, due to 1-based indexing */
-    memcpy( ctx->T_1_pub_key, &tree[1],
-            MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) );
+    memcpy( ctx->T_1_pub_key, &tree[node_bytes], node_bytes );
 
     ctx->have_public_key = 1;
 
     ret = 0;
 
 exit:
-    mbedtls_platform_zeroize( tree, sizeof( tree ) );
+    mbedtls_platform_zeroize( tree, node_bytes *
+        MERKLE_TREE_NODE_AM(priv_ctx->params.type) );
+    mbedtls_free ( tree );
 
     return( ret );
 }