diff --git a/snark/src/lib.rs b/snark/src/lib.rs
index c4a82f9..96e6cfe 100644
--- a/snark/src/lib.rs
+++ b/snark/src/lib.rs
@@ -31,6 +31,11 @@ extern "C" {
         bt1: *mut G1,
         bt2: *mut G2,
         ct: *mut G1);
+    fn libsnarkwrap_test_compare_key(
+        kp: *const libc::c_void,
+        size_of_queries: libc::uint64_t,
+        k_query: *const G1
+    ) -> bool;
     fn libsnarkwrap_test_keygen(
         cs: *const libc::c_void,
         tau: *const Fr,
@@ -109,6 +114,17 @@ impl Keypair {
             }
         }
     }
+
+    pub fn compare(&self, k_query: &[G1]) -> bool {
+
+        unsafe {
+            libsnarkwrap_test_compare_key(
+                self.ptr,
+                k_query.len() as u64,
+                &k_query[0]
+            )
+        }
+    }
 }
 
 impl CS {
diff --git a/snark/src/libsnarkwrap.cpp b/snark/src/libsnarkwrap.cpp
index 2fc23a5..a81e05d 100644
--- a/snark/src/libsnarkwrap.cpp
+++ b/snark/src/libsnarkwrap.cpp
@@ -277,6 +277,23 @@ extern "C" void* libsnarkwrap_test_keygen(
     );
 }
 
+extern "C" bool libsnarkwrap_test_compare_key(
+    const r1cs_ppzksnark_keypair<curve_pp> *kp,
+    uint64_t size_of_queries,
+    const curve_G1 *k_query
+)
+{
+    assert(kp->pk.K_query.size() == size_of_queries);
+
+    bool ret = true;
+
+    for (size_t i = 0; i < size_of_queries; i++) {
+        ret &= (k_query[i] == kp->pk.K_query[i]);
+    }
+
+    return ret;
+}
+
 extern "C" bool libsnarkwrap_test_compare_tau(
     const curve_G1 *inputs1,
     const curve_G2 *inputs2,
diff --git a/src/protocol.rs b/src/protocol.rs
index 42c202a..4aba5f5 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -536,10 +536,18 @@ fn implthing() {
     // Initializing pk_K as pk_A + pk _B + pk_C
     let mut pk_K = Vec::with_capacity(pk_A.len());
 
-    for ((&a, &b), &c) in pk_A.iter().zip(pk_B_temp.iter()).zip(pk_C.iter()) {
+    for ((&a, &b), &c) in pk_A.iter().take(pk_A.len() - 1)
+                              .zip(pk_B_temp.iter().take(pk_B.len() - 1))
+                              .zip(pk_C.iter().take(pk_C.len() - 1))
+    {
         pk_K.push(a + b + c);
     }
 
+    // Perform Z extention as libsnark does.
+    pk_K.push(pk_A[pk_A.len() - 1]);
+    pk_K.push(pk_B_temp[pk_B_temp.len() - 1]);
+    pk_K.push(pk_C[pk_C.len() - 1]);
+
     for (i, player) in players.iter().enumerate() {
         let (
             new_vk_gamma,
@@ -580,4 +588,6 @@ fn implthing() {
     }
 
     let target_kp = shared_secrets.keypair(&cs);
+
+    assert!(target_kp.compare(&pk_K));
 }