#include <fstream>
#include <iostream>
#include <sstream>
#include <vector>

using namespace std;

long checksum(vector<long> blocks) {
  long sum = 0;

  for (long i = 0; i < blocks.size(); i++) {
    if (blocks[i] < 0) continue;
    sum += i * blocks[i];
  }

  return sum;
}

vector<long> parseDisk(const char *path) {
  ifstream input(path);
  string line;
  getline(input, line);

  char c;
  stringstream ss;
  ss << line;
  bool swap = true;

  // cout << line << endl;

  vector<long> blocks;
  long id = 0;

  while (ss >> c) {
    int num = c - '0';
    if (swap) {
      for (int i = 0; i < num; i++) {
        blocks.push_back(id);
      }
      id++;
    } else {
      for (int i = 0; i < num; i++) {
        blocks.push_back(-1);
      }
    }
    swap = !swap;
  }
  return blocks;
}

void printDisk(vector<long> blocks) {
  for (int i = 0; i < blocks.size(); i++) {
    if (blocks[i] < 0) {
      cout << '.';
    } else {
      cout << blocks[i];
    }
  }
  cout << endl;
}

long part1(const char* path) {
  vector<long> blocks = parseDisk(path);

  while (true) {
    bool flag = false;
    for (int i = 0; i < blocks.size(); i++) {
      if (blocks[i] == -1) {
        blocks[i] = blocks[blocks.size() - 1];
        blocks.pop_back();
        flag = true;
        break;
      }
    }

    if (!flag) break;
  }

  return checksum(blocks);
}

long part2(const char *path) {
  vector<long> blocks = parseDisk(path);

  // printDisk(blocks);

  int current_size = 1;
  long current_id = blocks[blocks.size() - 1];
  for (int i = blocks.size() - 2; i >= 0; i--) {
    if (blocks[i] == current_id) {
      current_size++;
    } else {
      // skip empty block
      if (current_id == -1) {
        current_size = 1;
        current_id = blocks[i];
        continue;
      }

      // try finding a new spot
      int empty_size = 0;
      for (int j = 0; j <= i + 1; j++) {
        // found empty spot
        if (empty_size >= current_size) {
          // cout << "empty_size=" << empty_size << " current_size=" << current_size << endl;
          for (int k = 0; k < current_size; k++) {
            blocks[j - k - 1] = current_id;
            blocks[i + 1 + k] = -1;
          }
          // printDisk(blocks);
          break;
        }
        // check for next empty spot
        if (blocks[j] == -1) {
          empty_size++;
        } else {
          empty_size = 0;
        }
      }
      current_size = 1;
      current_id = blocks[i];
    }
  }

  return checksum(blocks);
}

int main(int argc, char *argv[]) {
  cout << "example1: " << part1("example") << endl;
  cout << "input1: " << part1("input") << endl;
  cout << "example2: " << part2("example") << endl;
  cout << "input2: " << part2("input") << endl;
  return 0;
}