Task Arrangement -- incomplete

Task schedule: Given a sequence of task like A B C, and a coldtime, which means you need to wait for that much time to start next[same] task. Now---

Input: string , n

Output: the best task finishing sequence.

Eg" Input: AAABBB, 2

Output: AB_AB_AB ("_" means do nothing and wait.) Source: 点击打开链接

Maybe, first see a problem easier than this one. (This version was terrible wrong!)

#include <string>
#include <iostream>
#include <unordered_map>
#include <vector>
using namespace std;

/*
  Give a task and some cool down time k (same task cool down distance).
  calcualte the total sum of time to finish the task.
*/
int totalTime(string str) {
  int count = str.size();
  unordered_map<char, int> charToIndex;
  for(int i = 0; i < str.size(); ++i) {
    char tmp = str[i];
    char index = i;
    if(charToIndex.find(tmp) == charToIndex.end()) {
      charToIndex.insert({tmp, index});
      continue;
    } else {
      if(index - charToIndex[tmp] > 2) {
        charToIndex[tmp] = index;
        continue;
      } else {
        if(index - charToIndex[tmp] == 1) count += 2;
        else count += 1;
        charToIndex[tmp] = index;
      }
    }
  }
  return count;
}

int main(void) {
  cout << totalTime("AABCB") << endl;
}

 The index need to be shifted once found one repeation.

// task cool down time k
#include "header.h"
using namespace std;

int taskCoolDown(vector<int>& nums, int k) {
  int totalTime = 0;
  int shift = 0;
  unordered_map<int, int> timeToIndex;
  for(int i = 0; i < nums.size(); ++i) {
    if(timeToIndex.find(nums[i]) == timeToIndex.end()) {
      timeToIndex.insert({nums[i], i + shift});
      totalTime++;
    } else {
      int tmp = timeToIndex[nums[i]];
      if(k + 1 - (i + shift - tmp) > 0) totalTime += k + 2 - (i + shift - tmp);
      else totalTime += 1;
      if(k + 1 - (i - tmp) > 0) shift += k + 1 - (i - tmp);
      else shift = 0;
      timeToIndex[nums[i]] = i + shift;
    }
  }
  return totalTime;
}

int main(void) {
  vector<int> nums{1, 1, 2, 1};
  vector<int> nums_1{1, 2, 3, 1, 2, 3};
  cout << taskCoolDown(nums, 2) << endl;
}