diff --git a/polar_list_decoder.hh b/polar_list_decoder.hh index 83ce286..56f41cf 100644 --- a/polar_list_decoder.hh +++ b/polar_list_decoder.hh @@ -55,24 +55,30 @@ struct PolarListNode TYPE sft = soft[1]; PATH fork[2*TYPE::SIZE]; for (int k = 0; k < TYPE::SIZE; ++k) - fork[k] = fork[k+TYPE::SIZE] = metric[k]; + fork[2*k] = fork[2*k+1] = metric[k]; for (int k = 0; k < TYPE::SIZE; ++k) if (sft.v[k] < 0) - fork[k] -= sft.v[k]; + fork[2*k] -= sft.v[k]; else - fork[k+TYPE::SIZE] += sft.v[k]; + fork[2*k+1] += sft.v[k]; int perm[2*TYPE::SIZE]; for (int k = 0; k < 2*TYPE::SIZE; ++k) perm[k] = k; - std::nth_element(perm, perm+TYPE::SIZE, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; }); +#if 1 + std::stable_sort(perm, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; }); +#else + std::partial_sort(perm, perm+TYPE::SIZE-1, perm+2*TYPE::SIZE, [fork](int a, int b){ + return fork[a] < fork[b] ? true : fork[a] > fork[b] ? false : a < b; + }); +#endif for (int k = 0; k < TYPE::SIZE; ++k) metric[k] = fork[perm[k]]; MAP map; for (int k = 0; k < TYPE::SIZE; ++k) - map.v[k] = perm[k] % TYPE::SIZE; + map.v[k] = perm[k] >> 1; TYPE hrd; for (int k = 0; k < TYPE::SIZE; ++k) - hrd.v[k] = perm[k] < TYPE::SIZE ? 1 : -1; + hrd.v[k] = 1 - 2 * (perm[k] & 1); message[*count] = hrd; maps[*count] = map; ++*count; diff --git a/polar_parity_aided.hh b/polar_parity_aided.hh index f913d95..4068075 100644 --- a/polar_parity_aided.hh +++ b/polar_parity_aided.hh @@ -103,24 +103,30 @@ struct PolarParityNode TYPE sft = soft[1]; PATH fork[2*TYPE::SIZE]; for (int k = 0; k < TYPE::SIZE; ++k) - fork[k] = fork[k+TYPE::SIZE] = metric[k]; + fork[2*k] = fork[2*k+1] = metric[k]; for (int k = 0; k < TYPE::SIZE; ++k) if (sft.v[k] < 0) - fork[k] -= sft.v[k]; + fork[2*k] -= sft.v[k]; else - fork[k+TYPE::SIZE] += sft.v[k]; + fork[2*k+1] += sft.v[k]; int perm[2*TYPE::SIZE]; for (int k = 0; k < 2*TYPE::SIZE; ++k) perm[k] = k; - std::nth_element(perm, perm+TYPE::SIZE, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; }); +#if 1 + std::stable_sort(perm, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; }); +#else + std::partial_sort(perm, perm+TYPE::SIZE-1, perm+2*TYPE::SIZE, [fork](int a, int b){ + return fork[a] < fork[b] ? true : fork[a] > fork[b] ? false : a < b; + }); +#endif for (int k = 0; k < TYPE::SIZE; ++k) metric[k] = fork[perm[k]]; MAP map; for (int k = 0; k < TYPE::SIZE; ++k) - map.v[k] = perm[k] % TYPE::SIZE; + map.v[k] = perm[k] >> 1; TYPE hrd; for (int k = 0; k < TYPE::SIZE; ++k) - hrd.v[k] = perm[k] < TYPE::SIZE ? 1 : -1; + hrd.v[k] = 1 - 2 * (perm[k] & 1); if (*count) { message[*index] = hrd; *parity = PH::qmul(vshuf(*parity, map), hrd);