Как найти количество путей в дереве с определенным весом?

Question

Яндекс Контест. Задача G. Пути в дереве

Ссылка на задачу: https://contest.yandex.ru/contest/36783/problems/G

Т.к. для просмотра нужна регистрация, то для удобства напишу текст сюда

Дано укоренённое дерево на n вершинах и число X. В каждой вершине записано число — её вес

Назовём восходящим путь ai, p(a_i), p(p(a_i)), ..., где p(a) — родитель вершины a. Проще говоря, восходящий путь — это путь, который начинается с некоторой вершины и двигается в сторону корня (не обязательно доходя до него). Путь может состоять из одной вершины

Весом пути назовём суммарный вес вершин на этом пути.

Найдите количество восходящих путей с весом X

Ограничения:

1 ≤ n ≤ 2 * 10^5
−10^9 ≤ X ≤ 10^9
p_i - номер родителя вершины: 0 ≤ p_i ≤ n − 1, либо −1 если это корень
w_i - вес вершины: −10^4 ≤ w_i ≤ 10^4

Как решил я:

Начинаю идти с конца массива (дерева) и добавляю вес к сумме. Начинаю цикл, пока не дойду до корня и сумма не превосходит X, перехожу к родителю и повторяю предыдущее действие. После окончания цикла смотрю сумма в точности равна X или нет, если да, то увеличиваю счётчик на 1, иначе не делаю ничего. Перед переходом у следующей вершине обнуляю сумму

Код (интерактивный):

class Vertex {
    constructor(weight, parent) {
        this.weight = weight;
        this.parent = parent;
    }
}

function getNumberOfUpgoingPaths(tree, x) {
    let count = 0;
    let sum = 0;
    
    for (let i = tree.length - 1; i >= 0; --i) {
        const vertex = tree[i];
        let parentIndex = vertex.parent;
        
        sum += vertex.weight;
        
        while(parentIndex !== -1 && sum < x) {
            const parentVertex = tree[parentIndex];
            
            sum += parentVertex.weight;
            
            parentIndex = parentVertex.parent;
        }
        
        if (sum === x) ++count;
        
        sum = 0;
    }
    
    return count;
}

// Обработка входных данных

const inputs = document.querySelector('#inputs');

const startAlgorithm = () => {
  console.clear();
  
  const tree = [];
  const [[countOfvertexes, x], ...vertexes] = inputs.value.split('\n').map(x => x.split(' ').map(Number));
  
  vertexes.forEach(vertex => tree.push(new Vertex(vertex[1], vertex[0])))
  
  console.log(getNumberOfUpgoingPaths(tree, x));
}

startAlgorithm();

inputs.addEventListener('input', startAlgorithm);

textarea {
  width: 90%;
  height: 400px;
}

<textarea id="inputs">6 3
-1 1
0 1
0 1
1 1
2 2
3 1</textarea>

Если хотите менять данные, то формат ввода таков:

В первой строке первые 2 числа - это количество вершин и X (число вершин при обработке игнорируется, но писать его нужно, чтобы не ломать формат ввода данных)
Начиная со второй строки перечисляются вершины, где левое число - это номер родителя, а правое число - это вес вершины

Проблема:

На каком-то из тестов (в яндексе не показывается на каком тесте) он ломается - возвращает неверный результат. На тестах, которые я сам придумывал, алгоритм работает верно. Помогите пожалуйста найти ошибку в рассуждениях или хотябы тест кейс, на котором всё ломается, дальше уже думаю смогу сам исправить ошибку

Answer 1

Вы упустили, что вес вершины может быть отрицательным

w_i - вес вершины: −10^4 ≤ w_i ≤ 10^4

Поэтому && sum < x может остановить рано.

А по эффективности - в последней теме вы должны были кой-чему научиться.

Answer 2

короче решенная задача на с++, основанная на вот этом тексте: https://www.techiedelight.com/ru/count-paths-with-given-sum-binary-tree/ Но по-моему там ошибка, потому что там сначала прибавляется map[sum_so_far] += 1; и только потом считается long count = map[sum_so_far - k]; у меня этот алгоритм работал некорректно. Этот прошел тест

#include <vector>
#include <iostream>
#include <set>
#include <algorithm>
#include <map>
#include <unordered_map>

using namespace std;

struct Vertex
{
  long weight;
  long parent;

  Vertex(long w1, long p1)
  {
    weight = w1;
    parent = p1;
  }
};

vector<vector<long>> createReverseTree(const vector<Vertex>& treeOrig) {
  long n = treeOrig.size();
  vector<vector<long>> reverseTree(n);

  for (int i = 0; i < n; i++) {
    int parent = treeOrig[i].parent;

    if (parent < 0)
    {
      continue;
    }
    reverseTree[parent].push_back(i);
  }

  return reverseTree;
}


// Функция для рекурсивного обхода дерева от корня до листьев
long traverseTree(
  const vector<vector<long>>& reverseTree,
  vector<Vertex>& treeOrig,
  long node,
  long k,
  long sum_so_far,
  std::unordered_map<long, long>& map
) 
{
  Vertex& curVertex = treeOrig[node];
  sum_so_far += curVertex.weight;
  

  // получаем количество путей с суммой `k`, которые заканчиваются текущим узлом
  long count = map[sum_so_far - k];
  map[sum_so_far] += 1;

  for (long child : reverseTree[node]) 
  {
    count += traverseTree(reverseTree, treeOrig, child, k, sum_so_far, map);
  }

  map[sum_so_far] -= 1;

  return count;
}

int root = 0;

vector<Vertex> readTree(long n) 
{
  vector<Vertex> tree;

  for (long i = 0; i < n; i++)
  {
    long parent, weight;
    cin >> parent >> weight;
    Vertex v(weight, parent);

    if (parent < 0)
    {
      root = i;
    }
    tree.push_back(v);
  }
  return tree;
}

long getNumberOfUpgoingPaths(vector<Vertex>& treeOrig, long x)
{
  long resVal = 0;

  auto tree = createReverseTree(treeOrig);

  unordered_map<long, long> map;
  map[0] = 1;

  resVal = traverseTree(tree, treeOrig, root, x, 0, map);

  return resVal;
}



int main() {
  int n;
  cin >> n;
  long x;
  cin >> x;
  vector<Vertex> tree = readTree(n);
  cout << getNumberOfUpgoingPaths(tree, x);
}

БЛОГ НА HUSL

Как найти количество путей в дереве с определенным весом?

Ответы (2 шт):