diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index f001a8253..e7a47d0c1 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -688,12 +688,14 @@ static const int palette_color_context_lookup[PALETTE_COLOR_CONTEXTS] = {
 };
 
 int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
-                                  int c, int n, int *color_order) {
+                                  int c, int n, uint8_t *color_order,
+                                  int *color_idx) {
   int i, j, max, max_idx, temp;
   int scores[PALETTE_MAX_SIZE + 10];
   int weights[4] = { 3, 2, 3, 2 };
   int color_ctx = 0;
   int color_neighbors[4];
+  int inverse_color_order[PALETTE_MAX_SIZE];
   assert(n <= PALETTE_MAX_SIZE);
   if (c - 1 >= 0)
     color_neighbors[0] = color_map[r * cols + c - 1];
@@ -711,7 +713,10 @@ int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
     color_neighbors[3] = color_map[(r - 1) * cols + c + 1];
   else
     color_neighbors[3] = -1;
-  for (i = 0; i < PALETTE_MAX_SIZE; ++i) color_order[i] = i;
+  for (i = 0; i < PALETTE_MAX_SIZE; ++i) {
+    color_order[i] = i;
+    inverse_color_order[i] = i;
+  }
   memset(scores, 0, PALETTE_MAX_SIZE * sizeof(scores[0]));
   for (i = 0; i < 4; ++i) {
     if (color_neighbors[i] >= 0) scores[color_neighbors[i]] += weights[i];
@@ -734,6 +739,8 @@ int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
       temp = color_order[i];
       color_order[i] = color_order[max_idx];
       color_order[max_idx] = temp;
+      inverse_color_order[color_order[i]] = i;
+      inverse_color_order[color_order[max_idx]] = max_idx;
     }
   }
   for (i = 0; i < 4; ++i) color_ctx = color_ctx * 11 + scores[i];
@@ -743,6 +750,9 @@ int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
       break;
     }
   if (color_ctx >= PALETTE_COLOR_CONTEXTS) color_ctx = 0;
+  if (color_idx != NULL) {
+    *color_idx = inverse_color_order[color_map[r * cols + c]];
+  }
   return color_ctx;
 }
 #endif  // CONFIG_PALETTE
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index f46e879ca..871c2a42f 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -222,7 +222,8 @@ static INLINE int av1_ceil_log2(int n) {
 
 #if CONFIG_PALETTE
 int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
-                                  int c, int n, int *color_order);
+                                  int c, int n, uint8_t *color_order,
+                                  int *color_idx);
 #endif  // CONFIG_PALETTE
 
 #ifdef __cplusplus
diff --git a/av1/decoder/detokenize.c b/av1/decoder/detokenize.c
index 0abcbaaf6..e67cdeaa8 100644
--- a/av1/decoder/detokenize.c
+++ b/av1/decoder/detokenize.c
@@ -282,15 +282,15 @@ static int decode_coefs(const MACROBLOCKD *xd, PLANE_TYPE type,
 #if CONFIG_PALETTE
 void av1_decode_palette_tokens(MACROBLOCKD *const xd, int plane,
                                aom_reader *r) {
-  MODE_INFO *const mi = xd->mi[0];
-  MB_MODE_INFO *const mbmi = &mi->mbmi;
+  const MODE_INFO *const mi = xd->mi[0];
+  const MB_MODE_INFO *const mbmi = &mi->mbmi;
   const BLOCK_SIZE bsize = mbmi->sb_type;
   const int rows = (4 * num_4x4_blocks_high_lookup[bsize]) >>
                    (xd->plane[plane != 0].subsampling_y);
   const int cols = (4 * num_4x4_blocks_wide_lookup[bsize]) >>
                    (xd->plane[plane != 0].subsampling_x);
-  int color_idx, color_ctx, color_order[PALETTE_MAX_SIZE];
-  int n = mbmi->palette_mode_info.palette_size[plane != 0];
+  uint8_t color_order[PALETTE_MAX_SIZE];
+  const int n = mbmi->palette_mode_info.palette_size[plane != 0];
   int i, j;
   uint8_t *color_map = xd->plane[plane != 0].color_index_map;
   const aom_prob (*const prob)[PALETTE_COLOR_CONTEXTS][PALETTE_COLORS - 1] =
@@ -299,10 +299,10 @@ void av1_decode_palette_tokens(MACROBLOCKD *const xd, int plane,
 
   for (i = 0; i < rows; ++i) {
     for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
-      color_ctx =
-          av1_get_palette_color_context(color_map, cols, i, j, n, color_order);
-      color_idx = aom_read_tree(r, av1_palette_color_tree[n - 2],
-                                prob[n - 2][color_ctx], ACCT_STR);
+      const int color_ctx = av1_get_palette_color_context(color_map, cols, i, j,
+                                                          n, color_order, NULL);
+      const int color_idx = aom_read_tree(r, av1_palette_color_tree[n - 2],
+                                          prob[n - 2][color_ctx], ACCT_STR);
       assert(color_idx >= 0 && color_idx < n);
       color_map[i * cols + j] = color_order[color_idx];
     }
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index e7cb2141c..d863ea711 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -888,8 +888,7 @@ static int rd_pick_palette_intra_sby(
   if (colors > 1 && colors <= 64) {
     int r, c, i, j, k;
     const int max_itr = 50;
-    int color_ctx, color_idx = 0;
-    int color_order[PALETTE_MAX_SIZE];
+    uint8_t color_order[PALETTE_MAX_SIZE];
     float *const data = x->palette_buffer->kmeans_data_buf;
     float centroids[PALETTE_MAX_SIZE];
     uint8_t *const color_map = xd->plane[0].color_index_map;
@@ -970,13 +969,9 @@ static int rd_pick_palette_intra_sby(
               1);
       for (i = 0; i < rows; ++i) {
         for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
-          color_ctx = av1_get_palette_color_context(color_map, cols, i, j, k,
-                                                    color_order);
-          for (r = 0; r < k; ++r)
-            if (color_map[i * cols + j] == color_order[r]) {
-              color_idx = r;
-              break;
-            }
+          int color_idx;
+          const int color_ctx = av1_get_palette_color_context(
+              color_map, cols, i, j, k, color_order, &color_idx);
           assert(color_idx >= 0 && color_idx < k);
           this_rate += cpi->palette_y_color_cost[k - 2][color_ctx][color_idx];
         }
@@ -1785,8 +1780,7 @@ static void rd_pick_palette_intra_sbuv(
   if (colors > 1 && colors <= 64) {
     int r, c, n, i, j;
     const int max_itr = 50;
-    int color_ctx, color_idx = 0;
-    int color_order[PALETTE_MAX_SIZE];
+    uint8_t color_order[PALETTE_MAX_SIZE];
     int64_t this_sse;
     float lb_u, ub_u, val_u;
     float lb_v, ub_v, val_v;
@@ -1876,13 +1870,9 @@ static void rd_pick_palette_intra_sbuv(
 
       for (i = 0; i < rows; ++i) {
         for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
-          color_ctx = av1_get_palette_color_context(color_map, cols, i, j, n,
-                                                    color_order);
-          for (r = 0; r < n; ++r)
-            if (color_map[i * cols + j] == color_order[r]) {
-              color_idx = r;
-              break;
-            }
+          int color_idx;
+          const int color_ctx = av1_get_palette_color_context(
+              color_map, cols, i, j, n, color_order, &color_idx);
           assert(color_idx >= 0 && color_idx < n);
           this_rate += cpi->palette_uv_color_cost[n - 2][color_ctx][color_idx];
         }
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 969aa9b1a..fdb19c68a 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -341,16 +341,16 @@ static INLINE int get_tx_eob(const struct segmentation *seg, int segment_id,
 }
 
 #if CONFIG_PALETTE
-void av1_tokenize_palette_sb(struct ThreadData *const td, BLOCK_SIZE bsize,
-                             int plane, TOKENEXTRA **t) {
-  MACROBLOCK *const x = &td->mb;
-  MACROBLOCKD *const xd = &x->e_mbd;
-  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
-  uint8_t *color_map = xd->plane[plane != 0].color_index_map;
-  PALETTE_MODE_INFO *pmi = &mbmi->palette_mode_info;
-  int n = pmi->palette_size[plane != 0];
-  int i, j, k;
-  int color_new_idx = -1, color_ctx, color_order[PALETTE_MAX_SIZE];
+void av1_tokenize_palette_sb(const struct ThreadData *const td,
+                             BLOCK_SIZE bsize, int plane, TOKENEXTRA **t) {
+  const MACROBLOCK *const x = &td->mb;
+  const MACROBLOCKD *const xd = &x->e_mbd;
+  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
+  const uint8_t *const color_map = xd->plane[plane != 0].color_index_map;
+  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
+  const int n = pmi->palette_size[plane != 0];
+  int i, j;
+  uint8_t color_order[PALETTE_MAX_SIZE];
   const int rows = (4 * num_4x4_blocks_high_lookup[bsize]) >>
                    (xd->plane[plane != 0].subsampling_y);
   const int cols = (4 * num_4x4_blocks_wide_lookup[bsize]) >>
@@ -361,15 +361,10 @@ void av1_tokenize_palette_sb(struct ThreadData *const td, BLOCK_SIZE bsize,
 
   for (i = 0; i < rows; ++i) {
     for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
-      color_ctx =
-          av1_get_palette_color_context(color_map, cols, i, j, n, color_order);
-      for (k = 0; k < n; ++k)
-        if (color_map[i * cols + j] == color_order[k]) {
-          color_new_idx = k;
-          break;
-        }
+      int color_new_idx;
+      const int color_ctx = av1_get_palette_color_context(
+          color_map, cols, i, j, n, color_order, &color_new_idx);
       assert(color_new_idx >= 0 && color_new_idx < n);
-
       (*t)->token = color_new_idx;
       (*t)->context_tree = probs[n - 2][color_ctx];
       (*t)->skip_eob_node = 0;
diff --git a/av1/encoder/tokenize.h b/av1/encoder/tokenize.h
index ed1be4662..c6d1e6078 100644
--- a/av1/encoder/tokenize.h
+++ b/av1/encoder/tokenize.h
@@ -55,8 +55,8 @@ struct AV1_COMP;
 struct ThreadData;
 
 #if CONFIG_PALETTE
-void av1_tokenize_palette_sb(struct ThreadData *const td, BLOCK_SIZE bsize,
-                             int plane, TOKENEXTRA **t);
+void av1_tokenize_palette_sb(const struct ThreadData *const td,
+                             BLOCK_SIZE bsize, int plane, TOKENEXTRA **t);
 #endif  // CONFIG_PALETTE
 
 void av1_tokenize_sb(const struct AV1_COMP *cpi, struct ThreadData *td,