HackerRank Separate the Chocolate problem solution

In this HackerRank Separate the Chocolate problem solution Tom and Derpina have a rectangular-shaped chocolate bar with chocolates labeled T, D, and U. They want to split the bar into exactly two pieces such that:

1. Tom’s piece can not contain any chocolate labeled D and similarly, Derpina’s piece can not contain any chocolate labeled T and U can be used by either of the two.
2. All chocolates in each piece must be connected (two chocolates are connected if they share an edge), i.e. the chocolates should form one connected component
3. The absolute difference between the number of chocolates in pieces should be at most K

After dividing it into exactly two pieces, in any piece, there should not be 4 adjacent chocolates that form a square, i.e. there should not be a fragment like this:

XX

XX

Problem solution in Python2.

from collections import defaultdict

M, N, K = (int(s) for s in raw_input().split())

sa=[]
D=[]
T=[]
max_d=-1
max_t=-1
for i in range(M):
    sa.append(raw_input())
if (N>M):
    sa=map(''.join ,zip(*sa))
    N,M=M,N
if (N==1):
    sa=map(''.join ,zip(*sa))
    N,M=M,N
    
for i in range(M):       
    s=sa[i][:N]
    sd = s
    st = s
    sd = sd.replace("U","0")
    sd = sd.replace("D","1")
    sd = sd.replace("T","0")
    if int(sd,2)!=0:
        max_d=i
    D.append(int(sd,2))
    st = st.replace("U","0")
    st = st.replace("D","0")
    st = st.replace("T","1")
    if int(sd,2)!=0:
        max_t=i
    T.append(int(st,2))

if (N%2==0)and(M%2==0):
    max_K = (N+M-1)        
else:
    max_K = (N+M+1)
max_K = max(max_K,4)
do_K = (K<max_K)    
    
def transfer(prev_comb, comb_compact):
        temp_comb_compact=comb_compact
        comb=[]
        for _ in range(N):
            comb.append(2*(temp_comb_compact%2)-1)
            temp_comb_compact /=2

        if prev_comb:
            label0_prev=-min(min(prev_comb),0)
            label1_prev=max(max(prev_comb),0)
        else:     
            prev_comb=[0]*N
            label0_prev=0
            label1_prev=0

        N_full=N+label0_prev+label1_prev+1
        N_half=N+label0_prev
        label = [0]*N_full
        adj = [[] for _ in range(N_full)]
        for i in range(N-1):
            if comb[i]*comb[i+1]>0:
                adj[i].append(i+1)
                adj[i+1].append(i)
                if (prev_comb[i]*prev_comb[i+1]>0)and(comb[i]*prev_comb[i]>0):
                    return []

        for i in range(N):
            prev_level = prev_comb[i]
            if comb[i]*prev_level>0:
                adj[i].append(N_half+prev_level)
                adj[N_half+prev_level].append(i)
    
        plus_label=0
        minus_label=0

        for i in range(N):
            if not(label[i])and(comb[i]!=0):
                if comb[i]>0:
                    plus_label += 1
                    current_label = plus_label
                elif comb[i]<0:
                    minus_label -= 1
                    current_label =minus_label
                label[i] = current_label
            
                q=[]
                q.append(i)
                while q:
                    node = q.pop()
                    for node_neighbour in adj[node]:
                        if not label[node_neighbour]:
                            label[node_neighbour] = current_label
                            q.append(node_neighbour)

        if (label0_prev>0)and(max(label[N:N_half])==0)and((label0_prev>1)or((label0_prev==1)and(minus_label!=0))):
            return []
        if (label1_prev>0)and(min(label[N_half+1:N_full])==0)and((label1_prev>1)or((label1_prev==1)and(plus_label!=0))):        
            return []

        return (tuple(label[:N]))


transfer_matrix=[]
transfer_matrix_end=[]
label_to_index={}
number_of_comb=[]
end_labels =[]

def build_transfer_matrix():
    N_labels = 0       
    pow2N=pow(2,N)
    queue_to_process = []
    for comb in range(pow2N):        
        labeled_comb = transfer([], comb)
        if labeled_comb:
            label_to_index[labeled_comb] = N_labels
            if (max(labeled_comb)<=1)and(min(labeled_comb)>=-1):
                end_labels.append(N_labels)
            queue_to_process.append([labeled_comb, N_labels])
            N_labels+=1
            transfer_matrix.append([])
            transfer_matrix_end.append([])
            if (comb&T[0]!=0)or(~comb&D[0]!=0):
                number_of_comb.append(0)
            else:
                number_of_comb.append(1)

    while queue_to_process:
        prev = queue_to_process.pop()
        prev_label = prev[0] 
        prev_index = prev[1]

        for comb in range(pow2N):
        
            labeled_comb = transfer(prev_label, comb)
            if labeled_comb:
                if labeled_comb not in label_to_index:
                    label_to_index[labeled_comb] = N_labels
                    if (max(labeled_comb)<=1)and(min(labeled_comb)>=-1):
                        end_labels.append(N_labels)
                    queue_to_process.append([labeled_comb, N_labels])
                    next_index = N_labels
                    N_labels+=1
                    transfer_matrix.append([])
                    transfer_matrix_end.append([])
                    number_of_comb.append(0)
                else:
                    next_index = label_to_index[labeled_comb]
                if (comb == 0)or(comb==pow2N-1):
                    transfer_matrix_end[prev_index].append([next_index,comb])
                else:
                    transfer_matrix[prev_index].append([next_index,comb])

build_transfer_matrix()

if do_K:
  K_lt = [2*bin(i).count("1")-N for i in range(pow(2,N))]

  prev_number_of_comb= number_of_comb
  number_of_comb=[[0]*(2*max_K+1) for _ in range(len(number_of_comb))]
  for comb in range(pow(2,N)):
    number_of_comb[comb][K_lt[comb]+max_K]=prev_number_of_comb[comb]
  
  for row in range(1,M):
    prev_number_of_comb= number_of_comb
    number_of_comb=[[0]*(2*max_K+1) for _ in range(len(number_of_comb))]
    for prev_id in range(len(transfer_matrix)):
        for K_id in range(2*max_K+1):
            if prev_number_of_comb[prev_id][K_id]: 
                for next_id,next_comb in transfer_matrix[prev_id]:      
                    if (next_comb&T[row]==0)and(~next_comb&D[row]==0)and(abs(K_id-max_K+K_lt[next_comb])<=max_K):
                        number_of_comb[next_id][K_id+K_lt[next_comb]]+=prev_number_of_comb[prev_id][K_id]

    if row==M-1:  
        for prev_id in range(len(transfer_matrix_end)):
            for K_id in range(2*max_K+1):
                if transfer_matrix_end[prev_id]:
                    for next_id,next_comb in transfer_matrix_end[prev_id]:
                        if (next_comb&T[M-1]==0)and(~next_comb&D[M-1]==0) and (abs(K_id-max_K+K_lt[next_comb])<=max_K):
                            number_of_comb[next_id][K_id+K_lt[next_comb]]+=prev_number_of_comb[prev_id][K_id]

  final_sum=0
  for end_id in end_labels:
    for K_id in range(2*max_K+1):
        if abs(K_id-max_K)<=K:
            final_sum+=number_of_comb[end_id][K_id]
  print final_sum

else:
  for row in range(1,M):
    prev_number_of_comb= number_of_comb
    number_of_comb=[0]*len(number_of_comb)
    for prev_id in range(len(transfer_matrix)):
        if prev_number_of_comb[prev_id]: 
            for next_id,next_comb in transfer_matrix[prev_id]:      
                if (next_comb&T[row]==0)and(~next_comb&D[row]==0):
                    number_of_comb[next_id]+=prev_number_of_comb[prev_id]

    if row==M-1:  
        for prev_id in range(len(transfer_matrix_end)):
            if transfer_matrix_end[prev_id]:
                for next_id,next_comb in transfer_matrix_end[prev_id]:
                    if (next_comb&T[M-1]==0)and(~next_comb&D[M-1]==0):
                        number_of_comb[next_id]+=prev_number_of_comb[prev_id]

  final_sum=0
  for end_id in end_labels:
    final_sum+=number_of_comb[end_id]
  print final_sum

Problem solution in Java.

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Scanner;
import java.util.Set;

public class SeparateChoco {

    static int T = 0;
    static int D = 1;
    static int U = -1;
    // yes, Go language style...
    static class MapIntLong extends HashMap<Integer, Long> {
    }

    static class MapStringMapIntLong extends HashMap<String, MapIntLong> {
    }

    static class MapStringRow extends HashMap<String, Row> {

    }

    static class MapStringSetString extends HashMap<String, Set<String>> {

    }

    static class Piece {
        int id;
        int groupIdx;

        public Piece(int piece, int groupIndex) {
            this.id = piece;
            this.groupIdx = groupIndex;
        }

        public String toString() {
            if (id == 0) {
                return "" + (char) ('a' + groupIdx);
            } else {
                return "" + (char) ('A' + groupIdx);
            }
        }

        public boolean equals(Object o) {
            Piece other = (Piece) o;
            return id == other.id && groupIdx == other.groupIdx;
        }

        public int hashCode() {
            return 31 * id + groupIdx;
        }

        public Piece copy() {
            return new Piece(id, groupIdx);
        }
    }

    static class Row {
        Piece[] pieces;
        boolean isHiding = false;

        public Row(Piece[] pieces) {
            this.pieces = pieces.clone();
        }

        public Row(int template, int width) {
            pieces = new Piece[width];
            for (int i = 0; i < width; i++) {
                pieces[i] = new Piece(template % 2, -1);
                template /= 2;
            }
        }

        public boolean hasMatch(int[] constraint) {
            for (int i = 0; i < pieces.length; i++) {
                if (constraint[i] >= 0 && constraint[i] != pieces[i].id) {
                    return false;
                }
            }
            return true;
        }

        public int countOnes() {
            int c = 0;
            for (int i = 0; i < pieces.length; i++) {
                c += pieces[i].id;
            }
            return c;
        }

        public boolean isOfTwoPieces() {
            if (isHiding && isUniform())
                return true;
            for (int i = 0; i < pieces.length; i++) {
                if (pieces[i].groupIdx > 0)
                    return false;
            }
            return true;
        }

        public String toString() {
            StringBuilder sb = new StringBuilder();
            for (Piece c : pieces) {
                sb.append(c);
            }
            if (isHiding) {
                sb.append("!");
            }
            return sb.toString();
        }

        public boolean equals(Object o) {
            Row other = (Row) o;
            return toString().equals(other.toString());
        }

        public Row copy() {
            Piece[] pcs = new Piece[pieces.length];
            for (int i = 0; i < pcs.length; i++) {
                pcs[i] = pieces[i].copy();
            }
            Row copy = new Row(pcs);
            copy.isHiding = isHiding;
            return copy;
        }

        boolean isUniform() {
            int p = pieces[0].id;
            for (int i = 1; i < pieces.length; i++) {
                if (pieces[i].id != p)
                    return false;
            }
            return true;
        }

        public void validate() {
            int groupOne = 20;
            int groupTwo = 20;

            for (int i = 0; i < pieces.length; i++) {
                int curGroup = pieces[i].groupIdx;
                if (curGroup >= 0) {
                    continue;
                }
                int p = pieces[i].id;
                if (p == 0) {
                    for (int j = i; j < pieces.length && pieces[j].id == p && pieces[j].groupIdx < 0; j++) {
                        pieces[j].groupIdx = groupOne;
                    }
                    groupOne++;
                } else {
                    for (int j = i; j < pieces.length && pieces[j].id == p && pieces[j].groupIdx < 0; j++) {
                        pieces[j].groupIdx = groupTwo;
                    }
                    groupTwo++;
                }
            }

            int[] movesOne = new int[pieces.length + 20];
            int[] movesTwo = new int[pieces.length + 20];
            for (int i = 0; i < movesOne.length; i++) {
                movesOne[i] = -1;
                movesTwo[i] = -1;
            }
            groupOne = 0;
            groupTwo = 0;

            for (int i = 0; i < pieces.length; i++) {
                int curGroup = pieces[i].groupIdx;
                if (curGroup < 0) {
                    continue;
                }
                if (pieces[i].id == 0) {
                    if (movesOne[curGroup] >= 0) {
                        pieces[i].groupIdx = movesOne[curGroup];
                    } else {
                        movesOne[curGroup] = groupOne;
                        pieces[i].groupIdx = groupOne;
                        groupOne++;
                    }
                } else {
                    if (movesTwo[curGroup] >= 0) {
                        pieces[i].groupIdx = movesTwo[curGroup];
                    } else {
                        movesTwo[curGroup] = groupTwo;
                        pieces[i].groupIdx = groupTwo;
                        groupTwo++;
                    }
                }
            }

        }
    }

    public void setSize(int width, int height, int diff) {
        this.width = width;
        this.height = height;
        this.diff = diff;
    }

    public void setConstraint(int[][] constraint) {
        this.constraint = constraint;
    }

    public List<Row> createRow() {
        List<Row> rows = new ArrayList<>();
        int cnt = 0;
        Piece[] pcs = new Piece[width];
        int[] solution = new int[width + 1];
        int[] solutionsCount = new int[width + 1];
        Piece[][] solutions = new Piece[width + 1][width + 1];
        solutions[0][0] = new Piece(0, 0);
        solutions[0][1] = new Piece(1, 0);
        solutionsCount[0] = 2;
        solution[cnt] = -1;
        while (true) {
            while (solution[cnt] == solutionsCount[cnt] - 1) {
                if (cnt == 0) {
                    return rows;
                }
                cnt--;
            }
            solution[cnt]++;
            pcs[cnt] = solutions[cnt][solution[cnt]];
            cnt++;
            solutionsCount[cnt] = 0;
            if (cnt < width) {
                int curPc = pcs[cnt - 1].id;
                solutions[cnt][0] = new Piece(curPc, pcs[cnt - 1].groupIdx); //same
                solutionsCount[cnt]++;
                List<Piece> pcsStack = new ArrayList<>();
                int grpId = -1;
                for (int i = 0; i < cnt; i++) {
                    int j = pcsStack.indexOf(pcs[i]);
                    if (j >= 0) {
                        while (pcsStack.size() > j + 1) {
                            pcsStack.remove(pcsStack.size() - 1);
                        }
                    } else {
                        pcsStack.add(pcs[i]);
                        if (pcs[i].id == 1 - curPc) {
                            if (pcs[i].groupIdx > grpId) {
                                grpId = pcs[i].groupIdx;
                            }
                        }
                    }
                }
                for (Piece c : pcsStack) {
                    if (c.id == 1 - curPc) {
                        solutions[cnt][solutionsCount[cnt]] = c;
                        solutionsCount[cnt]++;
                    }
                }
                solutions[cnt][solutionsCount[cnt]] = new Piece(1 - curPc, grpId + 1);
                solutionsCount[cnt]++;
            } else {
                rows.add(new Row(pcs));
            }
            solution[cnt] = -1;
        }
    }

    public Row move(Row from, Row to) {
        if (from.isHiding) {
            if (width == 1 && from.pieces[0].id == to.pieces[0].id) {
                to = to.copy();
                to.isHiding = true;
                to.validate();
                return to;
            }
            return null;
        }
        for (int i = 0; i < from.pieces.length - 1; i++) {
            int p = from.pieces[i].id;
            if (p == from.pieces[i + 1].id && p == to.pieces[i].id && p == to.pieces[i + 1].id) {
                return null;
            }
        }
        from = from.copy();
        to = to.copy();
        for (int i = 0; i < from.pieces.length; i++) {
            to.pieces[i].groupIdx = -1;
        }
        for (int i = 0; i < from.pieces.length; i++) {
            int p = from.pieces[i].id;
            int grpId1 = from.pieces[i].groupIdx;
            int grpId2 = to.pieces[i].groupIdx;
            if (p == to.pieces[i].id && grpId1 != grpId2) {
                if (grpId2 >= 0) {
                    for (int j = 0; j < from.pieces.length; j++) {
                        int ga = (grpId1 < grpId2) ? grpId1 : grpId2;
                        int gb = (grpId1 > grpId2) ? grpId1 : grpId2;
                        if (p == from.pieces[j].id && gb == from.pieces[j].groupIdx) {
                            from.pieces[j].groupIdx = ga;
                        }
                        if (p == to.pieces[j].id && gb == to.pieces[j].groupIdx) {
                            to.pieces[j].groupIdx = ga;
                        }
                    }
                } else {
                    to.pieces[i].groupIdx = grpId1;
                    int j = i + 1;
                    while (j < from.pieces.length && to.pieces[j].id == p) {
                        to.pieces[j].groupIdx = grpId1;
                        j++;
                    }
                    j = i - 1;
                    while (j >= 0 && to.pieces[j].id == p) {
                        to.pieces[j].groupIdx = grpId1;
                        j--;
                    }
                }
            }
        }
        Set<Piece> accounted = new HashSet<>();
        for (int i = 0; i < from.pieces.length; i++) {
            if (from.pieces[i].id == to.pieces[i].id) {
                accounted.add(from.pieces[i]);
            }
        }
        Set<Piece> unaccounted = new HashSet<>();
        for (int i = 0; i < from.pieces.length; i++) {
            if (!accounted.contains(from.pieces[i]) && !unaccounted.contains(from.pieces[i])) {
                unaccounted.add(from.pieces[i]);
            }
        }
        if (unaccounted.size() > 1) {
            return null;
        }
        to.validate();
        if (unaccounted.size() == 1) {
            if (!to.isUniform()) {
                return null;
            } else {
                to.isHiding = true;
            }
        }
        return to;
    }

    public void build() {
        int powOfTwo = 1;
        for (int i = 0; i < width; i++) {
            powOfTwo *= 2;
        }
        List<Row> rows = createRow();
        for (Row row : rows) {
            this.rows.put(row.toString(), row);
            if (row.isUniform()) {
                Row s1 = row.copy();
                s1.isHiding = true;
                this.rows.put(s1.toString(), s1);
            }
        }
        for (Row s : this.rows.values()) {
            Set<String> ts = new HashSet<>();
            for (int i = 0; i < powOfTwo; i++) {
                Row t = new Row(i, width);
                t = move(s, t);
                if (t != null)
                    ts.add(t.toString());
            }
            moves.put(s.toString(), ts);
        }
        for (int i = 0; i < powOfTwo; i++) {
            Row newRow = new Row(i, width);
            newRow.validate();
            if (!newRow.hasMatch(constraint[0])) {
                continue;
            }
            MapIntLong v = new MapIntLong();
            v.put(newRow.countOnes(), 1l);
            counts.put(newRow.toString(), v);
        }
        for (int i = 0; i < height - 1; i++) {
            counts = addRow(counts, constraint[i + 1]);
        }
        long sum = sum(counts, diff, width * height);
        System.out.println(sum);
    }

    MapStringMapIntLong addRow(MapStringMapIntLong counts, int[] cs) {
        MapStringMapIntLong next = new MapStringMapIntLong();
        for (String s : counts.keySet()) {
            MapIntLong vs = counts.get(s);
            for (String n : moves.get(s)) {
                Row t = rows.get(n);
                if (!t.hasMatch(cs)) {
                    continue;
                }
                int dk = t.countOnes();
                if (!next.containsKey(n)) {
                    MapIntLong v = new MapIntLong();
                    for (int k : vs.keySet()) {
                        v.put(k + dk, vs.get(k));
                    }
                    next.put(n, v);
                } else {
                    MapIntLong v = next.get(n);
                    for (int k : vs.keySet()) {
                        if (!v.containsKey(k + dk)) {
                            v.put(k + dk, vs.get(k));
                        } else {
                            v.put(k + dk, v.get(k + dk) + vs.get(k));
                        }
                    }
                }
            }
        }
        return next;
    }

    long sum(MapStringMapIntLong counts, int diff, int size) {
        long result = 0;
        for (String s : counts.keySet()) {
            Row state = rows.get(s);
            if (state.isOfTwoPieces()) {
                for (int k : counts.get(s).keySet()) {
                    int k1 = size - k;
                    if (Math.abs(k - k1) <= diff) {
                        long d = counts.get(s).get(k);
                        result += d;
                    }
                }
            }
        }
        return result;
    }

    int width;
    int height;
    int diff;
    int[][] constraint;

    MapStringRow rows = new MapStringRow();
    MapStringSetString moves = new MapStringSetString();
    MapStringMapIntLong counts = new MapStringMapIntLong();

    public void run() {
        Scanner in = new Scanner(System.in);
        int m = in.nextInt(), n = in.nextInt(), k = in.nextInt();
        if (m == 0 || n == 0) {
            System.out.println(1);
            return;
        }
        setSize(n, m, k);
        int[][] cs = new int[m][n];
        for (int i = 0; i < m; i++) {
            String s = in.next();
            for (int j = 0; j < n; j++) {
                char ch = s.charAt(j);
                cs[i][j] = ch == 'T' ? T : ch == 'D' ? D : U;
            }
        }
        setConstraint(cs);
        build();
    }

    public static void main(String[] args) {
        new SeparateChoco().run();
    }
}

Problem solution in C++.

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <map>
#include <tuple>
using namespace std;

typedef tuple<int, int, int> tiii;
#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define ROF(i, a, b) for (int i = (b); --i >= (a); )
#define fi first
#define se second

int ri()
{
  int x;
  scanf("%d", &x);
  return x;
}

const int N = 8;
char a[N][N+1];

int main()
{
  int m = ri(), n = ri(), diff = ri(), rest[3] = {0, 0, n*m};
  long ans = 0;
  REP(i, m) {
    scanf("%s", a[i]);
    REP(j, n) {
      if (a[i][j] == 'D')
        rest[0]++;
      if (a[i][j] == 'T')
        rest[1]++;
    }
  }
  map<tuple<vector<int>, int, int>, long> cur, pre;
  vector<int> src(n+1);
  int srcy = 0;
  iota(src.begin(), src.end(), 0);
  for (int i = 1; i < n+1; i += 2)
    srcy |= 1 << i;
  cur.emplace(make_tuple(src,srcy,0), 1);
  REP(i, m)
    REP(j, n) {
      pre.swap(cur);
      cur.clear();
      for (auto &it: pre) {
        vector<int> L(n+1), R;
        int y, num;
        tie(R, y, num) = it.fi;
        REP(i, n+1)
          L[R[i]] = i;
        REP(d, 2) 
          if (a[i][j] != (d ? 'D' : 'T') && (! i || ! j || (y>>j-1 & 7) != (d ? 7 : 0))) { // not the other color, not square
            
            if (! i || (y>>j+1 & 1) == d || R[j+1] != j && R[j+1] != j+1) {
              vector<int> LL = L, RR = R;
              int yy = y;

              
              LL[RR[j]] = LL[j];
              RR[LL[j]] = RR[j];
              LL[j] = RR[j] = j;
              
              if (j && (y>>j-1 & 1) == d) {
                
                LL[RR[j] = RR[j-1]] = j;
                RR[LL[j] = j-1] = j;
              }
              
              if ((y>>j+1 & 1) == d) {
                
                int jj = RR[j];
                RR[LL[jj] = LL[j+1]] = jj;
                LL[RR[j] = j+1] = j;
              }
              
             
              if (d)
                yy |= 1 << j;
              else
                yy &= ~ (1 << j);
              if (j == n-1) {
               
                if (RR[n] != n)
                  LL[RR[n]] = LL[n], RR[LL[n]] = RR[n];
                
                ROF(i, 0, n)
                  LL[i+1] = LL[i]+1, RR[i+1] = RR[i]+1;
                LL[0] = RR[0] = 0;
                yy = (yy & (1<<n)-1) << 1;
                if (yy & 2)
                  yy &= ~ 1;
                else
                  yy |= 1;
              }
              cur[make_tuple(RR, yy, num+d)] += it.se;
            } else if (! rest[d^1] && (n == 1 || i == m-1)) {
              bool ok = true;
              
              FOR(k, j+1, n)
                if ((y>>k & 3) == (d ? 3 : 0))
                  ok = false;
              
              REP(k, j)
                if ((y>>k & 1) != d)
                  ok = false;
              FOR(k, j+2, n+1)
                if ((y>>k & 1) != d)
                  ok = false;
              if (ok && abs(2*(num+(d?rest[2]:0))-n*m) <= diff)
                ans += it.se;
            }
          }
      }

      if (a[i][j] == 'D')
        rest[0]--;
      if (a[i][j] == 'T')
        rest[1]--;
      rest[2]--;
  }
  for (auto &it: cur) {
    vector<int> R;
    int y, num;
    tie(R, y, num) = it.fi;
    if (abs(2*num-n*m) <= diff) {
      int comp = 0;
      FOR(j, 1, n+1)
        if (R[j] <= j)
          comp++;
      if (comp <= 2)
        ans += it.se;
    }
  }
  printf("%ldn", ans);
}