diff --git a/c/tests/test_stats.c b/c/tests/test_stats.c
index aa2a0a1df2..28fa88f5ac 100644
--- a/c/tests/test_stats.c
+++ b/c/tests/test_stats.c
@@ -1082,6 +1082,16 @@ test_single_tree_divergence_matrix(void)
     ret = tsk_treeseq_divergence_matrix(
         &ts, 2, sample_set_sizes, NULL, 0, NULL, TSK_STAT_BRANCH, result);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    sample_set_sizes[0] = 3;
+    sample_set_sizes[1] = 1;
+    ret = tsk_treeseq_divergence_matrix(
+        &ts, 2, sample_set_sizes, NULL, 0, NULL, TSK_STAT_BRANCH, result);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = tsk_treeseq_divergence_matrix(
+        &ts, 2, sample_set_sizes, NULL, 0, NULL, TSK_STAT_SITE, result);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
     /* assert_arrays_almost_equal(4, result, D_site); */
 
     verify_divergence_matrix(&ts, TSK_STAT_BRANCH);
diff --git a/c/tskit/trees.c b/c/tskit/trees.c
index 9dfbce05fb..a725dcee60 100644
--- a/c/tskit/trees.c
+++ b/c/tskit/trees.c
@@ -6768,10 +6768,12 @@ get_sample_set_index_map(const tsk_size_t num_nodes, const tsk_size_t num_sample
     for (j = 0; j < num_nodes; j++) {
         node_index_map[j] = TSK_NULL;
     }
+    /* printf("START\n"); */
     i = 0;
     for (j = 0; j < num_sample_sets; j++) {
         total_samples += sample_set_sizes[j];
         for (k = 0; k < sample_set_sizes[j]; k++) {
+            /* printf("j = %d k = %d i = %d\n", (int) j, (int) k, (int) i); */
             u = sample_sets[i];
             i++;
             if (u < 0 || u >= (tsk_id_t) num_nodes) {
@@ -6861,7 +6863,9 @@ tsk_treeseq_divergence_matrix(const tsk_treeseq_t *self, tsk_size_t num_sample_s
     N = num_sample_sets_in;
     if (sample_sets_in == NULL) {
         sample_sets = self->samples;
-        N = self->num_samples;
+        if (sample_set_sizes_in == NULL) {
+            N = self->num_samples;
+        }
     }
     sample_set_sizes = sample_set_sizes_in;
     /* If sample_set_sizes is NULL, assume its N 1S */
@@ -6876,6 +6880,7 @@ tsk_treeseq_divergence_matrix(const tsk_treeseq_t *self, tsk_size_t num_sample_s
         }
         sample_set_sizes = tmp_sample_set_sizes;
     }
+    /* printf("N = %d\n", (int) N); */
     ret = get_sample_set_index_map(num_nodes, N, sample_set_sizes, sample_sets,
         &total_samples, &sample_set_index_map);
     if (ret != 0) {